Add files using upload-large-folder tool

BIO/sft/qwen-production-08022302/v0-20250802-230250/checkpoint-1029-merged/added_tokens.json

@@ -0,0 +1,24 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

BioReason-main/data/README.md

@@ -0,0 +1,35 @@
+# BioReasoning Data Curation
+
+Jupyter notebooks for processing genetic variant data and creating ML datasets for biological reasoning tasks.
+
+## Notebooks
+
+**Core Analysis**
+- `BioReasoning_DataCuration_KEGG.ipynb` - KEGG pathway analysis with Claude API
+- `Clinvar_Coding.ipynb` - ClinVar variant processing and gene mapping
+- `Clinvar_SNV_Non_SNV.ipynb` - SNV/structural variant datasets with VEP annotations
+
+**KEGG Pipeline**  
+- `KEGG_Data_1.ipynb` - KEGG network data processing and variant identification
+- `KEGG_Data_2.ipynb` - Variant parsing and sequence generation
+- `KEGG_Data_3.ipynb` - Final ML dataset creation with Q&A pairs
+
+**Variant Prediction**
+- `VEP.ipynb` - Variant effect prediction datasets (ClinVar, OMIM, eQTL)
+
+## Setup
+
+```bash
+brew install brewsci/bio/edirect  # For ClinVar (macOS)
+export ANTHROPIC_API_KEY="your-key"  # For KEGG analysis
+```
+
+## Usage
+
+Each notebook has a configuration section - update paths/keys as needed, then run sequentially.
+
+**Key Outputs:**
+- KEGG biological reasoning datasets
+- ClinVar variant-disease associations  
+- VEP prediction task datasets
+- Genomic sequences with variant context

BioReason-main/grpo_trainer_lora_model/adapter_config.json

@@ -0,0 +1,37 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "unsloth/qwen2.5-1.5b-instruct-unsloth-bnb-4bit",
+  "bias": "none",
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": false,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 64,
+  "lora_bias": false,
+  "lora_dropout": 0,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 64,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "o_proj",
+    "gate_proj",
+    "v_proj",
+    "up_proj",
+    "q_proj",
+    "down_proj",
+    "k_proj"
+  ],
+  "task_type": "CAUSAL_LM",
+  "use_dora": false,
+  "use_rslora": false
+}

BioReason-main/grpo_trainer_lora_model/ds_config_stage2.json

@@ -0,0 +1,41 @@
+{
+  "bf16": {
+    "enabled": true
+  },
+  "optimizer": {
+    "type": "AdamW",
+    "params": {
+      "lr": "auto",
+      "betas": "auto",
+      "eps": "auto",
+      "weight_decay": "auto"
+    }
+  },
+  "scheduler": {
+    "type": "WarmupLR",
+    "params": {
+      "warmup_min_lr": "auto",
+      "warmup_max_lr": "auto",
+      "warmup_num_steps": "auto"
+    }
+  },
+  "zero_optimization": {
+    "stage": 2,
+    "offload_optimizer": {
+      "device": "cpu",
+      "pin_memory": true
+    },
+    "contiguous_gradients": true,
+    "overlap_comm": true,
+    "allgather_partitions": true,
+    "allgather_bucket_size": 5e8,
+    "reduce_scatter": true,
+    "reduce_bucket_size": 5e8
+  },
+  "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
+}

BioReason_new/bioreason/dataset/protein.py

@@ -0,0 +1,421 @@
+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.pl.processing_pl import ProteinLLMProcessor
+from trl.data_utils import maybe_apply_chat_template
+
+
+class ProteinDataset(Dataset):
+    """Dataset for protein-text paired 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)
+
+            with open(file_path, "r", encoding="utf-8") as f:
+                items = json.load(f)
+                if isinstance(items, list):
+                    for item in items:
+                        processed_item = self._process_item(item)
+                        self.data.append(processed_item)
+                else:
+                    processed_item = self._process_item(items)
+                    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", [])
+        if isinstance(reasoning_steps, list):
+            reasoning = "\n".join(reasoning_steps)
+        else:
+            reasoning = str(reasoning_steps)
+
+        # Process protein sequence - remove any whitespace and convert to uppercase
+        protein_sequence = item.get("protein_sequence", "").replace(" ", "").upper().strip()
+        
+        # Handle protein description/function
+        protein_description = item.get("protein_description", item.get("function", "")).strip()
+
+        return {
+            "question": question,
+            "answer": answer,
+            "reasoning": reasoning,
+            "protein_sequence": protein_sequence,
+            "protein_description": protein_description,
+        }
+
+    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_protein_dataset(
+    dataset: ProteinDataset,
+    train_ratio: float = 0.8,
+    val_ratio: float = 0.1,
+    test_ratio: float = 0.1,
+    seed: int = 42,
+) -> Tuple[ProteinDataset, ProteinDataset, ProteinDataset]:
+    """
+    Split a protein 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
+        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_protein_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 protein 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 protein dataset
+    """
+    dataset = ProteinDataset(data_dir)
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+    )
+
+
+def get_format_protein_function(model_name: str) -> Any:
+    """
+    Get the appropriate format function for a given model name.
+    """
+    if model_name.lower() == "llm":
+        return format_protein_for_llm
+    elif model_name.lower() == "protein-llm":
+        return format_protein_for_protein_llm
+    else:
+        raise ValueError(f"Unsupported model name: {model_name}")
+
+
+def format_protein_for_protein_llm(example: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Format a protein example into the required chat format for Protein-LLM.
+    """
+    return {
+        "prompt": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "protein", "text": None},
+                    {"type": "text", "text": example["question"].strip()},
+                ],
+            },
+            {
+                "role": "assistant",
+                "reasoning_content": example["reasoning"].strip(),
+                "content": [
+                    {"type": "text", "text": f"Answer: {example['answer'].strip()}"},
+                ],
+            },
+        ],
+        "protein_sequences": [
+            example["protein_sequence"],
+        ],
+        "answer": example["answer"],
+    }
+
+
+def format_protein_for_llm(example: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Format a protein example into the required chat format for LLM.
+    """
+    question = f"Protein sequence: {example['protein_sequence']}\nQuestion: {example['question']}"
+    return {
+        "prompt": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "protein", "text": None},
+                    {"type": "text", "text": question.strip()},
+                ],
+            },
+            {
+                "role": "assistant",
+                "reasoning_content": example["reasoning"].strip(),
+                "content": [
+                    {"type": "text", "text": f"Answer: {example['answer'].strip()}"},
+                ],
+            },
+        ],
+        "protein_sequences": [
+            "",
+        ],
+        "answer": example["answer"],
+    }
+
+
+def format_protein_contrastive(example: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Format a protein example for contrastive learning.
+    """
+    # return {
+    #     "protein": example["protein"],
+    #     "text": example["text"],
+    # }
+    protein_seq = example.get("protein_sequence") or example.get("protein") or ""
+    text_desc = (example.get("protein_description") or 
+                example.get("text") or 
+                example.get("description") or 
+                example.get("function") or "")
+    
+    return {
+        "protein": protein_seq,
+        "text": text_desc,
+        "protein_sequence": protein_seq,  # 保持向后兼容
+        "text_description": text_desc,   # 保持向后兼容
+    }
+
+def protein_text_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, List[str]]:
+    """
+    修复后的 collate function for protein-text contrastive learning.
+    """
+    protein_sequences = []
+    text_sequences = []
+    
+    for item in batch:
+        # 尝试多个可能的字段名
+        protein_seq = (item.get("protein_sequence") or 
+                      item.get("protein") or "")
+        text_seq = (item.get("text_description") or 
+                   item.get("text") or 
+                   item.get("description") or "")
+        
+        protein_sequences.append(protein_seq)
+        text_sequences.append(text_seq)
+    
+    return {
+        "protein_sequences": protein_sequences,
+        "text_sequences": text_sequences,
+    }
+
+
+def protein_llm_collate_fn(
+    examples: List[Dict],
+    processor: ProteinLLMProcessor,
+    max_length_text: int,
+    max_length_protein: int,
+    return_answer_in_batch: bool = False,
+) -> Dict:
+    """
+    Custom collate function for Protein-LLM 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_protein_sequences = [example["protein_sequences"] for example in examples]
+
+    batch = processor(
+        text=prompts_text,
+        batch_protein_sequences=batch_protein_sequences,
+        return_tensors="pt",
+        padding=True,
+        padding_side="left",
+        add_special_tokens=False,
+        max_length_text=max_length_text,
+        max_length_protein=max_length_protein,
+    )
+
+    # 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 protein_collate_fn(
+    batch: List[Dict[str, Any]],
+    protein_tokenizer: Any,
+    label2id: Dict[str, int],
+    max_length: int = 1024,
+) -> Dict[str, Any]:
+    """
+    Custom collate function for protein models.
+    """
+    protein_sequences = [item["protein_sequence"] for item in batch]
+
+    # Tokenize protein sequences
+    tokenized_protein = protein_tokenizer(
+        protein_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 = {
+        "protein_ids": tokenized_protein.input_ids,
+        "protein_attention_mask": tokenized_protein.attention_mask,
+        "labels": labels_tensor,
+    }
+
+    return tokenized_batch

BioReason_new/bioreason/dataset/utils.py

@@ -0,0 +1,135 @@
+from datasets import Dataset as HfDataset
+import torch
+from torch.utils.data import Dataset
+from typing import Dict, Any
+
+
+def torch_to_hf_dataset(torch_dataset: Dataset) -> HfDataset:
+    """
+    Convert a PyTorch Dataset to a Hugging Face Dataset.
+    
+    Args:
+        torch_dataset: PyTorch Dataset to convert
+        
+    Returns:
+        HfDataset: Converted Hugging Face Dataset
+    """
+    # Extract all data from PyTorch dataset
+    data = []
+    for i in range(len(torch_dataset)):
+        data.append(torch_dataset[i])
+    
+    return HfDataset.from_list(data)
+
+
+def truncate_protein(example: Dict[str, Any], truncate_protein_per_side: int = 1024) -> Dict[str, Any]:
+    """
+    Truncate protein sequences to a maximum length.
+    
+    Args:
+        example: Dataset example containing protein sequences
+        truncate_protein_per_side: Maximum length to keep from each side
+        
+    Returns:
+        Dict[str, Any]: Modified example with truncated protein sequences
+    """
+    if "protein_sequence" in example:
+        protein_seq = example["protein_sequence"]
+        if len(protein_seq) > 2 * truncate_protein_per_side:
+            # Keep the first and last parts of the sequence
+            truncated_seq = protein_seq[:truncate_protein_per_side] + protein_seq[-truncate_protein_per_side:]
+            example["protein_sequence"] = truncated_seq
+            
+    if "protein_sequences" in example:
+        truncated_sequences = []
+        for seq in example["protein_sequences"]:
+            if len(seq) > 2 * truncate_protein_per_side:
+                truncated_seq = seq[:truncate_protein_per_side] + seq[-truncate_protein_per_side:]
+                truncated_sequences.append(truncated_seq)
+            else:
+                truncated_sequences.append(seq)
+        example["protein_sequences"] = truncated_sequences
+    
+    return example
+
+
+def clean_protein_sequence(sequence: str) -> str:
+    """
+    Clean protein sequence by removing invalid characters and normalizing.
+    
+    Args:
+        sequence: Raw protein sequence
+        
+    Returns:
+        str: Cleaned protein sequence
+    """
+    # Standard amino acid codes
+    valid_amino_acids = set("ACDEFGHIKLMNPQRSTVWY")
+    
+    # Remove whitespace and convert to uppercase
+    sequence = sequence.upper().replace(" ", "").replace("\n", "").replace("\t", "")
+    
+    # Keep only valid amino acid characters
+    cleaned_sequence = "".join([char for char in sequence if char in valid_amino_acids])
+    
+    return cleaned_sequence
+
+
+def validate_protein_example(example: Dict[str, Any]) -> bool:
+    """
+    Validate that a protein example has required fields and valid data.
+    
+    Args:
+        example: Dataset example to validate
+        
+    Returns:
+        bool: True if example is valid, False otherwise
+    """
+    # Check for required fields
+    required_fields = ["protein_sequence"]
+    for field in required_fields:
+        if field not in example or not example[field]:
+            return False
+    
+    # Check protein sequence validity
+    protein_seq = example["protein_sequence"]
+    if not isinstance(protein_seq, str) or len(protein_seq.strip()) == 0:
+        return False
+    
+    # Check for minimum sequence length (e.g., at least 10 amino acids)
+    cleaned_seq = clean_protein_sequence(protein_seq)
+    if len(cleaned_seq) < 10:
+        return False
+    
+    return True
+
+
+def format_protein_qa_example(example: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Format a protein example for question-answering tasks.
+    
+    Args:
+        example: Raw protein example
+        
+    Returns:
+        Dict[str, Any]: Formatted example
+    """
+    # Clean protein sequence
+    if "protein_sequence" in example:
+        example["protein_sequence"] = clean_protein_sequence(example["protein_sequence"])
+    
+    # Ensure answer is properly formatted
+    if "answer" in example:
+        answer = example["answer"]
+        if isinstance(answer, str):
+            example["answer"] = answer.strip().lower()
+        else:
+            example["answer"] = str(answer).strip().lower()
+    
+    # Format question if needed
+    if "question" in example:
+        question = example["question"]
+        if not question.endswith("?"):
+            example["question"] = question.strip() + "?"
+    
+    return example

BioReason_new/bioreason/models/pl/processing_pl.py

@@ -0,0 +1,279 @@
+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 ProteinLLMKwargs(CommonKwargs):
+    """Keyword arguments specific to protein 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: ProteinLLMKwargs
+    _defaults = {
+        "text_kwargs": {
+            "padding": False,
+        },
+    }
+
+class ProteinLLMProcessor(ProcessorMixin):
+    r"""
+    Constructs a ProteinLLM processor which wraps a ESM2 protein processor and a Qwen 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 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",
+    )
+    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 from ESM2)
+            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,
+        )
+            
+    # Move tensors to the specified device
+        if return_tensors == "pt" and device is not None:
+            protein_tokenized = {k: v.to(device) if isinstance(v, torch.Tensor) else v 
+                               for k, v in protein_tokenized.items()}
+            
+        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]:
+                    num_protein_tokens = (protein_processing_result['protein_tokenized']['input_ids'][index] != self.protein_tokenizer.pad_token_id).sum().item()
+                    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,
+        )
+        # Move text tensors to device if specified
+        if return_tensors == "pt" and device is not None:
+            text_inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v 
+                          for k, v in text_inputs.items()}
+        
+        # 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_new/bioreason/models/protein_llm.py

@@ -0,0 +1,1093 @@
+# import torch
+# import torch.nn as nn
+# from typing import Optional, List, Dict, Any, Union, Tuple
+# from transformers import (
+#     AutoTokenizer,
+#     AutoModelForCausalLM,
+#     EsmModel,
+#     EsmTokenizer,
+#     BertModel,
+#     BertTokenizer,
+# )
+
+# from bioreason.models.pl.processing_pl import ProteinLLMProcessor
+# #from bioreason.models.dl.chat_template_dl import CHAT_TEMPLATE
+
+
+
+import torch
+import torch.nn as nn
+from typing import Optional, List, Dict, Any, Union, Tuple
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    EsmModel,
+    EsmTokenizer,
+    BertModel,
+    BertTokenizer,
+    BertConfig,
+)
+
+from bioreason.models.pl.processing_pl import ProteinLLMProcessor
+#from bioreason.models.dl.chat_template_dl import CHAT_TEMPLATE
+
+
+class QFormerProjector(nn.Module):
+    """
+    QFormer-based projector that maps protein embeddings to text space.
+    Uses cross-attention mechanism for better alignment.
+    """
+    
+    import torch
+import torch.nn as nn
+from typing import Optional, List, Dict, Any, Union, Tuple
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    EsmModel,
+    EsmTokenizer,
+    BertModel,
+    BertTokenizer,
+)
+
+from bioreason.models.pl.processing_pl import ProteinLLMProcessor
+#from bioreason.models.dl.chat_template_dl import CHAT_TEMPLATE
+
+
+class QFormerProjector(nn.Module):
+    """
+    QFormer-based projector that maps protein embeddings to text space.
+    Uses cross-attention mechanism for better alignment.
+    """
+    
+    def __init__(
+        self,
+        protein_hidden_size: int,
+        text_hidden_size: int,
+        qformer_model_name: str = "microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext",
+        num_query_tokens: int = 32,
+        cross_attention_layers: int = 6,
+        max_protein_length: int = 400,  # Conservative limit: 32 + 400 = 432 < 512
+    ):
+        super().__init__()
+        
+        self.protein_hidden_size = protein_hidden_size
+        self.text_hidden_size = text_hidden_size
+        self.num_query_tokens = num_query_tokens
+        self.max_protein_length = max_protein_length  # 32 + 400 = 432 < 512 (BERT limit)
+        
+        # Load QFormer (BERT-based) - keep original config to avoid size mismatch
+        self.qformer = BertModel.from_pretrained(qformer_model_name)
+        self.qformer_hidden_size = self.qformer.config.hidden_size
+        
+        # Learnable query tokens
+        self.query_tokens = nn.Parameter(
+            torch.zeros(1, num_query_tokens, self.qformer_hidden_size)
+        )
+        self.query_tokens.data.normal_(mean=0.0, std=0.02)
+        
+        # Project protein features to QFormer dimension
+        self.protein_projection = nn.Linear(protein_hidden_size, self.qformer_hidden_size)
+        
+        # Final projection to text space
+        self.text_projection = nn.Linear(self.qformer_hidden_size, text_hidden_size)
+        
+        # Layer norm for stability
+        self.layer_norm = nn.LayerNorm(text_hidden_size)
+
+    def forward(
+        self,
+        protein_embeddings: torch.Tensor,  # [batch_size, seq_len, protein_hidden_size]
+        protein_attention_mask: torch.Tensor = None,  # [batch_size, seq_len]
+    ) -> torch.Tensor:
+        """
+        Forward pass through QFormer projector.
+        
+        Args:
+            protein_embeddings: Protein embeddings from ESM2
+            protein_attention_mask: Attention mask for protein sequences
+            
+        Returns:
+            Projected embeddings in text space [batch_size, num_query_tokens, text_hidden_size]
+        """
+        batch_size, seq_len, _ = protein_embeddings.size()
+        
+        # Truncate protein sequence if necessary
+        if seq_len > self.max_protein_length:
+            protein_embeddings = protein_embeddings[:, :self.max_protein_length, :]
+            if protein_attention_mask is not None:
+                protein_attention_mask = protein_attention_mask[:, :self.max_protein_length]
+            seq_len = self.max_protein_length
+        
+        # Project protein embeddings to QFormer dimension
+        protein_embeds = self.protein_projection(protein_embeddings)  # [B, L, H_qformer]
+        
+        # Expand query tokens for batch
+        query_tokens = self.query_tokens.expand(batch_size, -1, -1)  # [B, num_queries, H_qformer]
+        
+        # Create attention masks
+        query_atts = torch.ones(query_tokens.size()[:-1], dtype=torch.long, device=query_tokens.device)
+        
+        if protein_attention_mask is None:
+            protein_attention_mask = torch.ones(
+                protein_embeds.size()[:-1], dtype=torch.long, device=protein_embeds.device
+            )
+        
+        attention_mask = torch.cat([query_atts, protein_attention_mask], dim=1)
+        
+        # Ensure total length doesn't exceed model limit
+        total_length = attention_mask.size(1)
+        max_length = self.qformer.config.max_position_embeddings
+        
+        if total_length > max_length:
+            # Truncate protein sequence further if needed
+            excess = total_length - max_length
+            if excess > 0 and protein_embeds.size(1) > excess:
+                protein_embeds = protein_embeds[:, :-excess, :]
+                protein_attention_mask = protein_attention_mask[:, :-excess]
+                attention_mask = torch.cat([query_atts, protein_attention_mask], dim=1)
+            else:
+                raise ValueError(f"Cannot fit sequence into model max length {max_length}")
+        
+        # Combine embeddings
+        inputs_embeds = torch.cat([query_tokens, protein_embeds], dim=1)
+        
+        # Pass through QFormer without explicit position_ids (let model auto-generate)
+        outputs = self.qformer(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            return_dict=True,
+        )
+        
+        # Extract query outputs (first num_query_tokens)
+        query_output = outputs.last_hidden_state[:, :self.num_query_tokens, :]
+        
+        # Project to text space
+        text_embeds = self.text_projection(query_output)
+        text_embeds = self.layer_norm(text_embeds)
+        
+        return text_embeds
+
+
+class ProteinLLMModel(nn.Module):
+    """
+    A combined model that processes both protein sequences and text inputs.
+    Uses ESM2 for protein encoding, QFormer for projection, and Qwen for text generation.
+    """
+
+    def __init__(
+        self,
+        text_model_name: str,
+        protein_model_name: str,
+        qformer_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,
+        num_query_tokens: int = 32,
+        cross_attention_layers: int = 6,
+    ):
+        """
+        Initialize the ProteinLLMModel.
+
+        Args:
+            text_model_name: Name of the text model (Qwen)
+            protein_model_name: Name of the protein model (ESM2)
+            qformer_model_name: Name of the QFormer model
+            cache_dir: Directory to cache the models
+            max_length_protein: Maximum length of protein sequences
+            max_length_text: Maximum length of text sequences
+            text_model_finetune: Whether to finetune the text model
+            protein_model_finetune: Whether to finetune the protein model
+            num_query_tokens: Number of learnable query tokens
+            cross_attention_layers: Number of cross-attention layers in QFormer
+        """
+        super().__init__()
+
+        self.text_model_finetune = text_model_finetune
+        self.protein_model_finetune = protein_model_finetune
+        self.max_length_protein = max_length_protein
+        self.max_length_text = max_length_text
+        self.num_query_tokens = num_query_tokens
+
+        # Load the text model and tokenizer (Qwen)
+        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
+
+        # Add special tokens for protein
+        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|>")
+
+        # Load the protein model and tokenizer (ESM2)
+        self.protein_model = EsmModel.from_pretrained(
+            protein_model_name, cache_dir=cache_dir, trust_remote_code=True
+        )
+        self.protein_tokenizer = EsmTokenizer.from_pretrained(protein_model_name, trust_remote_code=True)
+        self.protein_config = self.protein_model.config
+
+        # Get model dimensions
+        self.text_hidden_size = self.text_config.hidden_size
+        self.protein_hidden_size = self.protein_config.hidden_size
+
+        # Create QFormer projector
+        self.protein_projection = QFormerProjector(
+            protein_hidden_size=self.protein_hidden_size,
+            text_hidden_size=self.text_hidden_size,
+            qformer_model_name=qformer_model_name,
+            num_query_tokens=num_query_tokens,
+            cross_attention_layers=cross_attention_layers,
+        )
+
+        # Create processor for handling inputs
+        self.processor = ProteinLLMProcessor(
+            tokenizer=self.text_tokenizer, 
+            protein_tokenizer=self.protein_tokenizer
+        )
+
+    def process_protein_embeddings(
+        self,
+        protein_tokenized: Dict[str, torch.Tensor],
+        batch_idx_map: List[int],
+        batch_size: int,
+    ) -> List[torch.Tensor]:
+        """
+        Process protein sequences to obtain embeddings.
+
+        Args:
+            protein_tokenized: Tokenized protein 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 protein representations
+        with torch.no_grad():
+            outputs = self.protein_model(
+                input_ids=protein_tokenized["input_ids"],
+                attention_mask=protein_tokenized["attention_mask"],
+            )
+            # Get the last hidden state
+            hidden_states = outputs.last_hidden_state  # shape: [n_seqs, seq_len, hidden_dim]
+
+        # Apply QFormer projection
+        hidden_states = hidden_states.to(
+            device=self.protein_projection.query_tokens.device, 
+            dtype=self.protein_projection.query_tokens.dtype
+        )
+        
+        # Project all embeddings at once
+        projected_states_list = []
+        for seq_idx in range(hidden_states.size(0)):
+            seq_embedding = hidden_states[seq_idx:seq_idx+1]  # [1, seq_len, hidden_dim]
+            seq_attention_mask = protein_tokenized["attention_mask"][seq_idx:seq_idx+1]
+            
+            projected_embedding = self.protein_projection(
+                seq_embedding, seq_attention_mask
+            )  # [1, num_query_tokens, text_hidden_size]
+            projected_states_list.append(projected_embedding.squeeze(0))  # [num_query_tokens, text_hidden_size]
+
+        # 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):
+            result[batch_idx].append(projected_states_list[seq_idx])
+
+        # 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), device=self.device)
+
+        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:
+        """
+        Forward pass through the model.
+        """
+        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 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()
+            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:
+                raise ValueError(
+                    f"Protein features and protein tokens do not match: features {n_protein_features}, tokens: {n_protein_tokens}"
+                )
+
+            # Ensure protein embeddings have the same dtype as the text embeddings
+            protein_embeds_flat = protein_embeds_flat.to(dtype=text_inputs_embeds.dtype)
+            text_inputs_embeds[mask] = protein_embeds_flat
+
+        # Forward pass through the text model
+        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]]:
+        """
+        Generate text based on protein and text inputs.
+        """
+        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 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()
+            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:
+                raise ValueError(
+                    f"Protein features and protein tokens do not match: features {n_protein_features}, tokens: {n_protein_tokens}"
+                )
+
+            # Ensure protein embeddings have the same dtype as the text embeddings
+            protein_embeds_flat = protein_embeds_flat.to(dtype=text_inputs_embeds.dtype)
+            text_inputs_embeds[mask] = protein_embeds_flat
+
+        # Generation
+        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
+    
+    @property
+    def device(self):
+        """Get the device of the model."""
+        return next(self.parameters()).device
+    
+    def to_device(self, tensor_dict):
+        """Move tensor dictionary to model device."""
+        return {k: v.to(self.device) if isinstance(v, torch.Tensor) else v 
+                for k, v in tensor_dict.items()}
+
+    def get_protein_embeddings(self, protein_sequences: List[str]) -> torch.Tensor:
+        """
+        Get raw protein embeddings before projection.
+        
+        Args:
+            protein_sequences: List of protein sequences
+            
+        Returns:
+            Raw protein embeddings [batch_size, seq_len, protein_hidden_size]
+        """
+        # Tokenize protein sequences
+        protein_inputs = self.protein_tokenizer(
+            protein_sequences,
+            padding=True,
+            truncation=True,
+            max_length=self.max_length_protein,
+            return_tensors="pt",
+        )
+
+        # Move to correct device
+        protein_inputs = self.to_device(protein_inputs)
+        
+        # Get protein embeddings
+        with torch.no_grad():
+            protein_outputs = self.protein_model(**protein_inputs)
+            protein_embeddings = protein_outputs.last_hidden_state
+        
+        return protein_embeddings
+
+    def get_protein_features(
+        self,
+        protein_sequences: List[str],
+        return_tensors: str = "pt",
+    ) -> torch.Tensor:
+        """
+        Extract protein features for contrastive learning.
+        
+        Args:
+            protein_sequences: List of protein sequences
+            return_tensors: Return format for tensors
+            
+        Returns:
+            Protein features [batch_size, num_query_tokens, text_hidden_size]
+        """
+        # Tokenize protein sequences
+        protein_inputs = self.protein_tokenizer(
+            protein_sequences,
+            padding=True,
+            truncation=True,
+            max_length=self.max_length_protein,
+            return_tensors=return_tensors,
+        )
+
+        protein_inputs = self.to_device(protein_inputs)
+        
+        # Get protein embeddings
+        with torch.no_grad():
+            protein_outputs = self.protein_model(**protein_inputs)
+            protein_embeddings = protein_outputs.last_hidden_state
+        
+        # Project through QFormer - Fixed: only pass two required arguments
+        protein_features = self.protein_projection(
+            protein_embeddings, protein_inputs["attention_mask"]
+        )
+        
+        # Global average pooling over query tokens
+      #  protein_features = protein_features.mean(dim=1)  # [batch_size, text_hidden_size]
+        
+        return protein_features
+
+    def get_text_features(
+        self,
+        text_sequences: List[str],
+        return_tensors: str = "pt",
+    ) -> torch.Tensor:
+        """
+        Extract text features for contrastive learning.
+        
+        Args:
+            text_sequences: List of text descriptions
+            return_tensors: Return format for tensors
+            
+        Returns:
+            Text features [batch_size, text_hidden_size]
+        """
+        # Tokenize text sequences
+        text_inputs = self.text_tokenizer(
+            text_sequences,
+            padding=True,
+            truncation=True,
+            max_length=self.max_length_text,
+            return_tensors=return_tensors,
+        )
+
+        text_inputs = self.to_device(text_inputs)
+        
+        # Get text embeddings from the embedding layer
+        with torch.no_grad():
+            text_embeddings = self.text_model.get_input_embeddings()(text_inputs["input_ids"])
+            
+            # Apply attention mask and average pooling
+            attention_mask = text_inputs["attention_mask"].unsqueeze(-1)
+            masked_embeddings = text_embeddings * attention_mask
+            text_features = masked_embeddings.sum(dim=1) / attention_mask.sum(dim=1)
+        
+        return text_features
+
+# class QFormerProjector(nn.Module):
+#     """
+#     QFormer-based projector that maps protein embeddings to text space.
+#     Uses cross-attention mechanism for better alignment.
+#     """
+    
+#     # def __init__(
+#     #     self,
+#     #     protein_hidden_size: int,
+#     #     text_hidden_size: int,
+#     #     qformer_model_name: str = "microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext",
+#     #     num_query_tokens: int = 32,
+#     #     cross_attention_layers: int = 6,
+#     # ):
+#     #     super().__init__()
+        
+#     #     self.protein_hidden_size = protein_hidden_size
+#     #     self.text_hidden_size = text_hidden_size
+#     #     self.num_query_tokens = num_query_tokens
+        
+#     #     # Load QFormer (BERT-based)
+#     #     self.qformer = BertModel.from_pretrained(qformer_model_name)
+#     #     self.qformer_hidden_size = self.qformer.config.hidden_size
+        
+#     #     # Learnable query tokens
+#     #     self.query_tokens = nn.Parameter(
+#     #         torch.zeros(1, num_query_tokens, self.qformer_hidden_size)
+#     #     )
+#     #     self.query_tokens.data.normal_(mean=0.0, std=0.02)
+        
+#     #     # Project protein features to QFormer dimension
+#     #     self.protein_projection = nn.Linear(protein_hidden_size, self.qformer_hidden_size)
+        
+#     #     # Final projection to text space
+#     #     self.text_projection = nn.Linear(self.qformer_hidden_size, text_hidden_size)
+        
+#     #     # Layer norm for stability
+#     #     self.layer_norm = nn.LayerNorm(text_hidden_size)
+        
+#     # def forward(
+#     #     self,
+#     #     protein_embeddings: torch.Tensor,  # [batch_size, seq_len, protein_hidden_size]
+#     #     protein_attention_mask: torch.Tensor = None,  # [batch_size, seq_len]
+#     # ) -> torch.Tensor:
+#     #     """
+#     #     Forward pass through QFormer projector.
+        
+#     #     Args:
+#     #         protein_embeddings: Protein embeddings from ESM2
+#     #         protein_attention_mask: Attention mask for protein sequences
+            
+#     #     Returns:
+#     #         Projected embeddings in text space [batch_size, num_query_tokens, text_hidden_size]
+#     #     """
+#     #     batch_size = protein_embeddings.size(0)
+        
+#     #     # Project protein embeddings to QFormer dimension
+#     #     protein_embeds = self.protein_projection(protein_embeddings)  # [B, L, H_qformer]
+        
+#     #     # Expand query tokens for batch
+#     #     query_tokens = self.query_tokens.expand(batch_size, -1, -1)  # [B, num_queries, H_qformer]
+        
+#     #     # Concatenate query tokens and protein embeddings
+#     #     query_atts = torch.ones(query_tokens.size()[:-1], dtype=torch.long, device=query_tokens.device)
+        
+#     #     if protein_attention_mask is None:
+#     #         protein_attention_mask = torch.ones(
+#     #             protein_embeds.size()[:-1], dtype=torch.long, device=protein_embeds.device
+#     #         )
+        
+#     #     attention_mask = torch.cat([query_atts, protein_attention_mask], dim=1)
+        
+#     #     # Create position ids
+#     #     position_ids = torch.arange(
+#     #         attention_mask.size(1), dtype=torch.long, device=attention_mask.device
+#     #     ).unsqueeze(0).expand(batch_size, -1)
+        
+#     #     # Combine embeddings
+#     #     inputs_embeds = torch.cat([query_tokens, protein_embeds], dim=1)
+        
+#     #     # Pass through QFormer
+#     #     outputs = self.qformer(
+#     #         inputs_embeds=inputs_embeds,
+#     #         attention_mask=attention_mask,
+#     #         position_ids=position_ids,
+#     #         return_dict=True,
+#     #     )
+        
+#     #     # Extract query outputs (first num_query_tokens)
+#     #     query_output = outputs.last_hidden_state[:, :self.num_query_tokens, :]
+        
+#     #     # Project to text space
+#     #     text_embeds = self.text_projection(query_output)
+#     #     text_embeds = self.layer_norm(text_embeds)
+        
+#     #     return text_embeds
+#     def __init__(
+#         self,
+#         protein_hidden_size: int,
+#         text_hidden_size: int,
+#         qformer_model_name: str = "microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext",
+#         num_query_tokens: int = 32,
+#         cross_attention_layers: int = 6,
+#         max_protein_length: int = 480,  # 新增：限制蛋白质序列长度
+#     ):
+#         super().__init__()
+        
+#         self.protein_hidden_size = protein_hidden_size
+#         self.text_hidden_size = text_hidden_size
+#         self.num_query_tokens = num_query_tokens
+#         self.max_protein_length = max_protein_length  # 32 + 480 = 512
+        
+#         # Load QFormer (BERT-based) with longer sequence support
+#         from transformers import BertConfig
+#         config = BertConfig.from_pretrained(qformer_model_name)
+        
+#         # 方案1：扩展模型的最大位置编码（如果原模型支持）
+#         config.max_position_embeddings = max(1024, num_query_tokens + max_protein_length)
+        
+#         self.qformer = BertModel.from_pretrained(qformer_model_name, config=config)
+#         self.qformer_hidden_size = self.qformer.config.hidden_size
+        
+#         # Learnable query tokens
+#         self.query_tokens = nn.Parameter(
+#             torch.zeros(1, num_query_tokens, self.qformer_hidden_size)
+#         )
+#         self.query_tokens.data.normal_(mean=0.0, std=0.02)
+        
+#         # Project protein features to QFormer dimension
+#         self.protein_projection = nn.Linear(protein_hidden_size, self.qformer_hidden_size)
+        
+#         # Final projection to text space
+#         self.text_projection = nn.Linear(self.qformer_hidden_size, text_hidden_size)
+        
+#         # Layer norm for stability
+#         self.layer_norm = nn.LayerNorm(text_hidden_size)
+
+#     def forward(
+#         self,
+#         protein_embeddings: torch.Tensor,  # [batch_size, seq_len, protein_hidden_size]
+#         protein_attention_mask: torch.Tensor = None,  # [batch_size, seq_len]
+#     ) -> torch.Tensor:
+#         """
+#         Forward pass through QFormer projector.
+        
+#         Args:
+#             protein_embeddings: Protein embeddings from ESM2
+#             protein_attention_mask: Attention mask for protein sequences
+            
+#         Returns:
+#             Projected embeddings in text space [batch_size, num_query_tokens, text_hidden_size]
+#         """
+#         batch_size, seq_len, _ = protein_embeddings.size()
+        
+#         # 方案2：截断蛋白质序列
+#         if seq_len > self.max_protein_length:
+#             protein_embeddings = protein_embeddings[:, :self.max_protein_length, :]
+#             if protein_attention_mask is not None:
+#                 protein_attention_mask = protein_attention_mask[:, :self.max_protein_length]
+#             seq_len = self.max_protein_length
+        
+#         # Project protein embeddings to QFormer dimension
+#         protein_embeds = self.protein_projection(protein_embeddings)  # [B, L, H_qformer]
+        
+#         # Expand query tokens for batch
+#         query_tokens = self.query_tokens.expand(batch_size, -1, -1)  # [B, num_queries, H_qformer]
+        
+#         # Create attention masks
+#         query_atts = torch.ones(query_tokens.size()[:-1], dtype=torch.long, device=query_tokens.device)
+        
+#         if protein_attention_mask is None:
+#             protein_attention_mask = torch.ones(
+#                 protein_embeds.size()[:-1], dtype=torch.long, device=protein_embeds.device
+#             )
+        
+#         attention_mask = torch.cat([query_atts, protein_attention_mask], dim=1)
+        
+#         # 确保总长度不超过模型限制
+#         total_length = attention_mask.size(1)
+#         max_length = self.qformer.config.max_position_embeddings
+        
+#         if total_length > max_length:
+#             raise ValueError(f"Total sequence length {total_length} exceeds model max length {max_length}")
+        
+#         # Combine embeddings
+#         inputs_embeds = torch.cat([query_tokens, protein_embeds], dim=1)
+        
+#         # 方案3：不使用position_ids，让模型自动生成
+#         # Pass through QFormer without explicit position_ids
+#         outputs = self.qformer(
+#             inputs_embeds=inputs_embeds,
+#             attention_mask=attention_mask,
+#             return_dict=True,
+#         )
+        
+#         # Extract query outputs (first num_query_tokens)
+#         query_output = outputs.last_hidden_state[:, :self.num_query_tokens, :]
+        
+#         # Project to text space
+#         text_embeds = self.text_projection(query_output)
+#         text_embeds = self.layer_norm(text_embeds)
+        
+#         return text_embeds
+
+
+# class ProteinLLMModel(nn.Module):
+#     """
+#     A combined model that processes both protein sequences and text inputs.
+#     Uses ESM2 for protein encoding, QFormer for projection, and Qwen for text generation.
+#     """
+
+#     def __init__(
+#         self,
+#         text_model_name: str,
+#         protein_model_name: str,
+#         qformer_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,
+#         num_query_tokens: int = 32,
+#         cross_attention_layers: int = 6,
+#     ):
+#         """
+#         Initialize the ProteinLLMModel.
+
+#         Args:
+#             text_model_name: Name of the text model (Qwen)
+#             protein_model_name: Name of the protein model (ESM2)
+#             qformer_model_name: Name of the QFormer model
+#             cache_dir: Directory to cache the models
+#             max_length_protein: Maximum length of protein sequences
+#             max_length_text: Maximum length of text sequences
+#             text_model_finetune: Whether to finetune the text model
+#             protein_model_finetune: Whether to finetune the protein model
+#             num_query_tokens: Number of learnable query tokens
+#             cross_attention_layers: Number of cross-attention layers in QFormer
+#         """
+#         super().__init__()
+
+#         self.text_model_finetune = text_model_finetune
+#         self.protein_model_finetune = protein_model_finetune
+#         self.max_length_protein = max_length_protein
+#         self.max_length_text = max_length_text
+#         self.num_query_tokens = num_query_tokens
+
+#         # Load the text model and tokenizer (Qwen)
+#         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
+
+#         # Add special tokens for protein
+#         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|>")
+
+#         # Load the protein model and tokenizer (ESM2)
+#         self.protein_model = EsmModel.from_pretrained(
+#             protein_model_name, cache_dir=cache_dir, trust_remote_code=True
+#         )
+#         self.protein_tokenizer = EsmTokenizer.from_pretrained(protein_model_name, trust_remote_code=True)
+#         self.protein_config = self.protein_model.config
+
+#         # Get model dimensions
+#         self.text_hidden_size = self.text_config.hidden_size
+#         self.protein_hidden_size = self.protein_config.hidden_size
+
+#         # Create QFormer projector
+#         self.protein_projection = QFormerProjector(
+#             protein_hidden_size=self.protein_hidden_size,
+#             text_hidden_size=self.text_hidden_size,
+#             qformer_model_name=qformer_model_name,
+#             num_query_tokens=num_query_tokens,
+#             cross_attention_layers=cross_attention_layers,
+#         )
+
+#         # Create processor for handling inputs
+#         self.processor = ProteinLLMProcessor(
+#             tokenizer=self.text_tokenizer, 
+#             protein_tokenizer=self.protein_tokenizer
+#         )
+
+#     def process_protein_embeddings(
+#         self,
+#         protein_tokenized: Dict[str, torch.Tensor],
+#         batch_idx_map: List[int],
+#         batch_size: int,
+#     ) -> List[torch.Tensor]:
+#         """
+#         Process protein sequences to obtain embeddings.
+
+#         Args:
+#             protein_tokenized: Tokenized protein 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 protein representations
+#         with torch.no_grad():
+#             outputs = self.protein_model(
+#                 input_ids=protein_tokenized["input_ids"],
+#                 attention_mask=protein_tokenized["attention_mask"],
+#             )
+#             # Get the last hidden state
+#             hidden_states = outputs.last_hidden_state  # shape: [n_seqs, seq_len, hidden_dim]
+
+#         # Apply QFormer projection
+#         hidden_states = hidden_states.to(
+#             device=self.protein_projection.query_tokens.device, 
+#             dtype=self.protein_projection.query_tokens.dtype
+#         )
+        
+#         # Project all embeddings at once
+#         projected_states_list = []
+#         for seq_idx in range(hidden_states.size(0)):
+#             seq_embedding = hidden_states[seq_idx:seq_idx+1]  # [1, seq_len, hidden_dim]
+#             seq_attention_mask = protein_tokenized["attention_mask"][seq_idx:seq_idx+1]
+            
+#             projected_embedding = self.protein_projection(
+#                 seq_embedding, seq_attention_mask
+#             )  # [1, num_query_tokens, text_hidden_size]
+#             projected_states_list.append(projected_embedding.squeeze(0))  # [num_query_tokens, text_hidden_size]
+
+#         # 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):
+#             result[batch_idx].append(projected_states_list[seq_idx])
+
+#         # 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,
+#         protein_tokenized: Optional[Dict[str, torch.Tensor]] = None,
+#         batch_idx_map: Optional[List[int]] = None,
+#         labels: Optional[torch.Tensor] = None,
+#         **kwargs,
+#     ) -> torch.Tensor:
+#         """
+#         Forward pass through the model.
+#         """
+#         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 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()
+#             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:
+#                 raise ValueError(
+#                     f"Protein features and protein tokens do not match: features {n_protein_features}, tokens: {n_protein_tokens}"
+#                 )
+
+#             # Ensure protein embeddings have the same dtype as the text embeddings
+#             protein_embeds_flat = protein_embeds_flat.to(dtype=text_inputs_embeds.dtype)
+#             text_inputs_embeds[mask] = protein_embeds_flat
+
+#         # Forward pass through the text model
+#         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]]:
+#         """
+#         Generate text based on protein and text inputs.
+#         """
+#         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 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()
+#             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:
+#                 raise ValueError(
+#                     f"Protein features and protein tokens do not match: features {n_protein_features}, tokens: {n_protein_tokens}"
+#                 )
+
+#             # Ensure protein embeddings have the same dtype as the text embeddings
+#             protein_embeds_flat = protein_embeds_flat.to(dtype=text_inputs_embeds.dtype)
+#             text_inputs_embeds[mask] = protein_embeds_flat
+
+#         # Generation
+#         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
+    
+
+#     @property
+#     def device(self):
+#         """Get the device of the model."""
+#         return next(self.parameters()).device
+    
+#     def to_device(self, tensor_dict):
+#         """Move tensor dictionary to model device."""
+#         return {k: v.to(self.device) if isinstance(v, torch.Tensor) else v 
+#                 for k, v in tensor_dict.items()}
+
+
+#     def get_protein_embeddings(self, protein_sequences: List[str]) -> torch.Tensor:
+#         """
+#         Get raw protein embeddings before projection.
+        
+#         Args:
+#             protein_sequences: List of protein sequences
+            
+#         Returns:
+#             Raw protein embeddings [batch_size, seq_len, protein_hidden_size]
+#         """
+#         # Tokenize protein sequences
+#         protein_inputs = self.protein_tokenizer(
+#             protein_sequences,
+#             padding=True,
+#             truncation=True,
+#             max_length=self.max_length_protein,
+#             return_tensors="pt",
+#         )
+
+#         # 移动到正确设备
+#         protein_inputs = self.to_device(protein_inputs)
+        
+#         # Get protein embeddings
+#         with torch.no_grad():
+#             protein_outputs = self.protein_model(**protein_inputs)
+#             protein_embeddings = protein_outputs.last_hidden_state
+        
+#         return protein_embeddings
+
+#     def get_protein_features(
+#         self,
+#         protein_sequences: List[str],
+#         return_tensors: str = "pt",
+#     ) -> torch.Tensor:
+#         """
+#         Extract protein features for contrastive learning.
+        
+#         Args:
+#             protein_sequences: List of protein sequences
+#             return_tensors: Return format for tensors
+            
+#         Returns:
+#             Protein features [batch_size, num_query_tokens, text_hidden_size]
+#         """
+#         # Tokenize protein sequences
+#         protein_inputs = self.protein_tokenizer(
+#             protein_sequences,
+#             padding=True,
+#             truncation=True,
+#             max_length=self.max_length_protein,
+#             return_tensors=return_tensors,
+#         )
+
+#         protein_inputs = self.to_device(protein_inputs)
+        
+#         # Get protein embeddings
+#         with torch.no_grad():
+#             protein_outputs = self.protein_model(**protein_inputs)
+#             protein_embeddings = protein_outputs.last_hidden_state
+        
+#         # Project through QFormer
+#         protein_features = self.protein_projection(
+#             protein_embeddings, protein_inputs["attention_mask"],token_type_ids=None
+#         )
+        
+#         # Global average pooling over query tokens
+#         protein_features = protein_features.mean(dim=1)  # [batch_size, text_hidden_size]
+        
+#         return protein_features
+
+#     def get_text_features(
+#         self,
+#         text_sequences: List[str],
+#         return_tensors: str = "pt",
+#     ) -> torch.Tensor:
+#         """
+#         Extract text features for contrastive learning.
+        
+#         Args:
+#             text_sequences: List of text descriptions
+#             return_tensors: Return format for tensors
+            
+#         Returns:
+#             Text features [batch_size, text_hidden_size]
+#         """
+#         # Tokenize text sequences
+#         text_inputs = self.text_tokenizer(
+#             text_sequences,
+#             padding=True,
+#             truncation=True,
+#             max_length=self.max_length_text,
+#             return_tensors=return_tensors,
+#         )
+
+#         text_inputs = self.to_device(text_inputs)
+        
+#         # Get text embeddings from the embedding layer
+#         with torch.no_grad():
+#             text_embeddings = self.text_model.get_input_embeddings()(text_inputs["input_ids"])
+            
+#             # Apply attention mask and average pooling
+#             attention_mask = text_inputs["attention_mask"].unsqueeze(-1)
+#             masked_embeddings = text_embeddings * attention_mask
+#             text_features = masked_embeddings.sum(dim=1) / attention_mask.sum(dim=1)
+        
+#         return text_features

BioReason_new/bioreason/protein_modules/_init_.py

@@ -0,0 +1,7 @@
+from .protein_base_module import ProteinBaseModule
+from .protein_module import ESM2ProteinModule
+
+__all__ = [
+    "ProteinBaseModule",
+    "ESM2ProteinModule",
+]

BioReason_new/bioreason/protein_modules/protein_base_module.py

@@ -0,0 +1,49 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Union
+import torch
+
+class ProteinBaseModule(ABC):
+    def __init__(self):
+        super().__init__()
+    
+    @abstractmethod
+    def get_proteinllm_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_proteinllm_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, proteins, return_tensors, padding, padding_side, add_special_tokens):
+        pass

BioReason_new/bioreason/protein_modules/protein_module.py

@@ -0,0 +1,257 @@
+from typing import Dict, Any, Union, List, Optional, Callable, Type
+from trl.data_utils import maybe_apply_chat_template
+import torch
+
+from bioreason.protein_modules.protein_module import ProteinBaseModule
+from bioreason.models.protein_llm import ProteinLLMModel
+from bioreason.models.dl.processing_dl import ProteinLLMProcessor
+
+
+class ESM2ProteinModule(ProteinBaseModule):
+    """
+    Protein module implementation for ESM2-based models with QFormer projection.
+
+    This module provides the interface between Protein-LLM models and the training
+    infrastructure, handling model loading, processing setup, and reward functions.
+    """
+
+    def __init__(self):
+        """Initialize the ESM2ProteinModule."""
+        super().__init__()
+
+    def get_proteinllm_key(self) -> str:
+        """
+        Get the key identifier for this Protein-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 "ProteinLLM" in model_id:
+            model_cls = ProteinLLMModel
+        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 Protein-LLM model.
+
+        Returns:
+            The processing class
+        """
+        return ProteinLLMProcessor
+
+    def get_proteinllm_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", "qformer", "projection"]
+
+    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 []
+
+    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")]
+
+    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_protein_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_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
+
+        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_protein = model.max_length_protein if not hasattr(model, 'module') else model.module.max_length_protein
+        
+        prompt_inputs = processing_class(
+            text=prompts_text,
+            batch_protein_sequences=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,
+        )
+
+        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 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 ESM2ProteinModule.iou_reward
+                case _:
+                    raise ValueError(f"Unsupported reward function: {func}")
+        elif func == "format":
+            match task_type:
+                case "rec":
+                    return ESM2ProteinModule.format_reward_rec
+                case _:
+                    raise ValueError(f"Unsupported reward function: {func}")
+        else:
+            raise ValueError(f"Unsupported reward function: {func}")

BioReason_new/bioreason/trainer/_init_.py

@@ -0,0 +1,11 @@
+from .grpo_config import DNALLMGRPOConfig
+from .grpo_trainer import DNALLMGRPOTrainer
+from .protein_grpo_config import ProteinLLMGRPOConfig
+from .protein_grpo_trainer import ProteinLLMGRPOTrainer
+
+__all__ = [
+    "DNALLMGRPOConfig",
+    "DNALLMGRPOTrainer", 
+    "ProteinLLMGRPOConfig",
+    "ProteinLLMGRPOTrainer",
+]

BioReason_new/bioreason/trainer/contrast_trainer.py

@@ -0,0 +1,372 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from typing import Dict, List, Optional, Any, Union
+from transformers import Trainer, TrainingArguments
+from dataclasses import dataclass, field
+import wandb
+from datasets import Dataset
+
+from bioreason.models.protein_llm import ProteinLLMModel
+
+def pl_concat_all_gather(tensor):
+    """
+    Gather tensors from all processes in distributed training.
+    Falls back to returning the original tensor if not in distributed mode.
+    """
+    if not dist.is_available() or not dist.is_initialized():
+        return tensor
+    
+    world_size = dist.get_world_size()
+    if world_size == 1:
+        return tensor
+    
+    gathered_tensors = [torch.zeros_like(tensor) for _ in range(world_size)]
+    dist.all_gather(gathered_tensors, tensor)
+    return torch.cat(gathered_tensors, dim=0)
+
+@dataclass
+class ContrastiveTrainingArguments(TrainingArguments):
+    """
+    Arguments for contrastive learning training.
+    """
+    temperature: float = field(
+        default=0.07,
+        metadata={"help": "Temperature parameter for contrastive loss"}
+    )
+    freeze_protein_model: bool = field(
+        default=True,
+        metadata={"help": "Whether to freeze the protein model during training"}
+    )
+    freeze_text_model: bool = field(
+        default=True,
+        metadata={"help": "Whether to freeze the text model during training"}
+    )
+    protein_weight: float = field(
+        default=1.0,
+        metadata={"help": "Weight for protein features in contrastive loss"}
+    )
+    text_weight: float = field(
+        default=1.0,
+        metadata={"help": "Weight for text features in contrastive loss"}
+    )
+    max_length_protein: int = field(
+        default=1024,
+        metadata={"help": "Maximum length for protein sequences"}
+    )
+    max_length_text: int = field(
+        default=512,
+        metadata={"help": "Maximum length for text sequences"}
+    )
+
+
+
+class ContrastiveLoss(nn.Module):
+    """
+    Contrastive loss for protein-text alignment.
+    """
+    
+    def __init__(self, temperature: float = 0.07):
+        super().__init__()
+        self.temperature = temperature
+        
+    def forward(
+        self, 
+        protein_features: torch.Tensor, 
+        text_features: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Compute contrastive loss between protein and text features.
+        
+        Args:
+            protein_features: [batch_size, hidden_size]
+            text_features: [batch_size, hidden_size]
+            
+        Returns:
+            Contrastive loss value
+        """
+        # Normalize features
+        protein_features = F.normalize(protein_features, dim=-1)
+        text_features = F.normalize(text_features, dim=-1)
+        
+        # Compute similarity matrix
+        similarity_matrix = torch.matmul(protein_features, text_features.T) / self.temperature
+        
+        # Create labels for positive pairs (diagonal elements)
+        batch_size = protein_features.size(0)
+        labels = torch.arange(batch_size, device=protein_features.device)
+        
+        # Compute cross-entropy loss in both directions
+        loss_p2t = F.cross_entropy(similarity_matrix, labels)
+        loss_t2p = F.cross_entropy(similarity_matrix.T, labels)
+        
+        # Average the losses
+        total_loss = (loss_p2t + loss_t2p) / 2
+        
+        return total_loss
+
+
+class ContrastiveTrainer(Trainer):
+    """
+    Trainer for contrastive learning between proteins and text.
+    """
+    
+    def __init__(
+        self,
+        model: ProteinLLMModel,
+        args: ContrastiveTrainingArguments,
+        train_dataset: Optional[Dataset] = None,
+        eval_dataset: Optional[Dataset] = None,
+        data_collator: Optional[callable] = None,
+        **kwargs
+    ):
+        self.contrastive_loss = ContrastiveLoss(temperature=args.temperature)
+        self.freeze_protein_model = args.freeze_protein_model
+        self.freeze_text_model = args.freeze_text_model
+        self.protein_weight = args.protein_weight
+        self.text_weight = args.text_weight
+        
+        # Freeze models if specified
+        if self.freeze_protein_model:
+            for param in model.protein_model.parameters():
+                param.requires_grad = False
+                
+        if self.freeze_text_model:
+            for param in model.text_model.parameters():
+                param.requires_grad = False
+        
+        # Ensure projection layers are trainable
+        for param in model.protein_projection.parameters():
+            param.requires_grad = True
+            
+        super().__init__(
+            model=model,
+            args=args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            data_collator=data_collator,
+            **kwargs
+        )
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        """
+        Compute contrastive loss.
+        
+        Args:
+            model: The ProteinLLMModel
+            inputs: Dictionary containing protein_sequences and text_sequences
+            return_outputs: Whether to return model outputs
+            
+        Returns:
+            Contrastive loss
+        """
+        protein_sequences = inputs["protein_sequences"]
+        text_sequences = inputs["text_sequences"]
+        
+        # Get protein features
+        protein_features = model.get_protein_features(protein_sequences)
+        
+        # Get text features  
+        text_features = model.get_text_features(text_sequences)
+        
+        # Compute contrastive loss
+        loss = self.contrastive_loss(protein_features, text_features)
+        
+        # Log metrics
+        with torch.no_grad():
+            # Compute similarity scores for monitoring
+            protein_features_norm = F.normalize(protein_features, dim=-1)
+            text_features_norm = F.normalize(text_features, dim=-1)
+            similarity_matrix = torch.matmul(protein_features_norm, text_features_norm.T)
+            
+            # Diagonal elements are positive pairs
+            positive_similarities = torch.diag(similarity_matrix)
+            negative_similarities = similarity_matrix[~torch.eye(similarity_matrix.size(0), dtype=bool)]
+            
+            self.log({
+                "contrastive_loss": loss.item(),
+                "positive_similarity_mean": positive_similarities.mean().item(),
+                "negative_similarity_mean": negative_similarities.mean().item(),
+                "positive_similarity_std": positive_similarities.std().item(),
+                "similarity_gap": (positive_similarities.mean() - negative_similarities.mean()).item(),
+            })
+        
+        if return_outputs:
+            outputs = {
+                "protein_features": protein_features,
+                "text_features": text_features,
+                "similarity_matrix": similarity_matrix,
+            }
+            return (loss, outputs)
+        
+        return loss
+
+    def evaluation_loop(self, dataloader, description, prediction_loss_only=None, ignore_keys=None, metric_key_prefix="eval"):
+        """
+        Custom evaluation loop for contrastive learning.
+        """
+        model = self._wrap_model(self.model, training=False, dataloader=dataloader)
+        model.eval()
+        
+        total_loss = 0.0
+        total_samples = 0
+        all_protein_features = []
+        all_text_features = []
+        
+        for step, inputs in enumerate(dataloader):
+            with torch.no_grad():
+                loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
+                
+                total_loss += loss.item()
+                total_samples += len(inputs["protein_sequences"])
+                
+                all_protein_features.append(outputs["protein_features"].cpu())
+                all_text_features.append(outputs["text_features"].cpu())
+        
+        # Compute overall metrics
+        avg_loss = total_loss / len(dataloader)
+        
+        # Concatenate all features
+        all_protein_features = torch.cat(all_protein_features, dim=0)
+        all_text_features = torch.cat(all_text_features, dim=0)
+        
+        # Compute retrieval metrics
+        retrieval_metrics = self.compute_retrieval_metrics(all_protein_features, all_text_features)
+        
+        metrics = {
+            f"{metric_key_prefix}_loss": avg_loss,
+            **{f"{metric_key_prefix}_{k}": v for k, v in retrieval_metrics.items()}
+        }
+        
+        return metrics
+
+    def compute_retrieval_metrics(self, protein_features: torch.Tensor, text_features: torch.Tensor) -> Dict[str, float]:
+        """
+        Compute retrieval metrics (Recall@K).
+        
+        Args:
+            protein_features: [num_samples, hidden_size]
+            text_features: [num_samples, hidden_size]
+            
+        Returns:
+            Dictionary of retrieval metrics
+        """
+        # Normalize features
+        protein_features = F.normalize(protein_features, dim=-1)
+        text_features = F.normalize(text_features, dim=-1)
+        
+        # Compute similarity matrix
+        similarity_matrix = torch.matmul(protein_features, text_features.T)
+        
+        # Protein-to-text retrieval
+        p2t_ranks = []
+        for i in range(similarity_matrix.size(0)):
+            similarities = similarity_matrix[i]
+            rank = (similarities >= similarities[i]).sum().item()
+            p2t_ranks.append(rank)
+        
+        # Text-to-protein retrieval  
+        t2p_ranks = []
+        for i in range(similarity_matrix.size(1)):
+            similarities = similarity_matrix[:, i]
+            rank = (similarities >= similarities[i]).sum().item()
+            t2p_ranks.append(rank)
+        
+        # Compute Recall@K
+        metrics = {}
+        for k in [1, 5, 10]:
+            p2t_recall_k = sum(1 for rank in p2t_ranks if rank <= k) / len(p2t_ranks)
+            t2p_recall_k = sum(1 for rank in t2p_ranks if rank <= k) / len(t2p_ranks)
+            
+            metrics[f"p2t_recall_at_{k}"] = p2t_recall_k
+            metrics[f"t2p_recall_at_{k}"] = t2p_recall_k
+            metrics[f"avg_recall_at_{k}"] = (p2t_recall_k + t2p_recall_k) / 2
+        
+        # Mean rank
+        metrics["p2t_mean_rank"] = sum(p2t_ranks) / len(p2t_ranks)
+        metrics["t2p_mean_rank"] = sum(t2p_ranks) / len(t2p_ranks)
+        
+        return metrics
+
+
+def protein_text_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, List[str]]:
+    """
+    Collate function for protein-text contrastive learning.
+    
+    Args:
+        batch: List of samples, each containing "protein_sequence" and "text_description"
+        
+    Returns:
+        Dictionary with lists of protein sequences and text descriptions
+    """
+    protein_sequences = [item["protein_sequence"] for item in batch]
+    text_sequences = [item["text_description"] for item in batch]
+    
+    return {
+        "protein_sequences": protein_sequences,
+        "text_sequences": text_sequences,
+    }
+
+
+# Example usage
+def train_contrastive_model(
+    model: ProteinLLMModel,
+    train_dataset: Dataset,
+    eval_dataset: Optional[Dataset] = None,
+    output_dir: str = "./contrastive_outputs",
+    num_epochs: int = 10,
+    batch_size: int = 32,
+    learning_rate: float = 1e-4,
+    temperature: float = 0.07,
+    **kwargs
+):
+    """
+    Train the model with contrastive learning.
+    
+    Args:
+        model: ProteinLLMModel to train
+        train_dataset: Training dataset with protein_sequence and text_description
+        eval_dataset: Optional evaluation dataset
+        output_dir: Directory to save outputs
+        num_epochs: Number of training epochs
+        batch_size: Training batch size
+        learning_rate: Learning rate
+        temperature: Temperature for contrastive loss
+        **kwargs: Additional training arguments
+    """
+    training_args = ContrastiveTrainingArguments(
+        output_dir=output_dir,
+        num_train_epochs=num_epochs,
+        per_device_train_batch_size=batch_size,
+        per_device_eval_batch_size=batch_size,
+        learning_rate=learning_rate,
+        temperature=temperature,
+        logging_steps=10,
+        evaluation_strategy="steps" if eval_dataset else "no",
+        eval_steps=100 if eval_dataset else None,
+        save_steps=500,
+        save_total_limit=3,
+        load_best_model_at_end=True if eval_dataset else False,
+        metric_for_best_model="eval_avg_recall_at_1" if eval_dataset else None,
+        greater_is_better=True,
+        report_to=["wandb"] if wandb.run else [],
+        **kwargs
+    )
+    
+    trainer = ContrastiveTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        data_collator=protein_text_collate_fn,
+    )
+    
+    # Train the model
+    trainer.train()
+    
+    # Save the final model
+    trainer.save_model()
+    
+    return trainer

BioReason_new/bioreason/trainer/contrast_trainer_new.py

@@ -0,0 +1,659 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from torch.utils.data import DataLoader
+from typing import Dict, List, Optional, Any, Union
+from dataclasses import dataclass, field
+from transformers import Trainer, TrainingArguments
+import wandb
+from datasets import Dataset
+
+from bioreason.models.protein_llm import ProteinLLMModel
+
+
+def pl_concat_all_gather(tensor):
+    """
+    Gather tensors from all processes in distributed training.
+    Falls back to returning the original tensor if not in distributed mode.
+    """
+    if not dist.is_available() or not dist.is_initialized():
+        return tensor
+    
+    world_size = dist.get_world_size()
+    if world_size == 1:
+        return tensor
+    
+    gathered_tensors = [torch.zeros_like(tensor) for _ in range(world_size)]
+    dist.all_gather(gathered_tensors, tensor)
+    return torch.cat(gathered_tensors, dim=0)
+
+
+@dataclass
+class ContrastiveTrainingArguments(TrainingArguments):
+    """
+    Arguments for contrastive learning training.
+    """
+    print(TrainingArguments.__module__)
+    print("----------")
+    temperature: float = field(
+        default=0.07,
+        metadata={"help": "Temperature parameter for contrastive loss"}
+    )
+    freeze_protein_model: bool = field(
+        default=True,
+        metadata={"help": "Whether to freeze the protein model during training"}
+    )
+    freeze_text_model: bool = field(
+        default=True,
+        metadata={"help": "Whether to freeze the text model during training"}
+    )
+    protein_weight: float = field(
+        default=1.0,
+        metadata={"help": "Weight for protein features in contrastive loss"}
+    )
+    text_weight: float = field(
+        default=1.0,
+        metadata={"help": "Weight for text features in contrastive loss"}
+    )
+    max_length_protein: int = field(
+        default=1024,
+        metadata={"help": "Maximum length for protein sequences"}
+    )
+    max_length_text: int = field(
+        default=512,
+        metadata={"help": "Maximum length for text sequences"}
+    )
+    enable_ptm: bool = field(
+        default=True,
+        metadata={"help": "Enable protein-text matching task"}
+    )
+    ptm_weight: float = field(
+        default=1.0,
+        metadata={"help": "Weight for protein-text matching loss"}
+    )
+
+
+class EnhancedContrastiveLoss(nn.Module):
+    """
+    Enhanced contrastive loss for protein-text alignment with multi-query support.
+    Based on BLIP2 QFormer implementation.
+    """
+    
+    def __init__(self, temperature: float = 0.07, enable_ptm: bool = True):
+        super().__init__()
+        self.temperature = temperature
+        self.enable_ptm = enable_ptm
+        if enable_ptm:
+            self.ptm_head = nn.Linear(768, 2)  # Assuming hidden size of 768
+        
+    def contrast_global(self, features_protein, features_text, features_protein_all, features_text_all, return_sim=False):
+        """
+        Compute global contrastive loss across all processes.
+        
+        Args:
+            features_protein: [B, num_queries, D] - local protein features
+            features_text: [B, D] - local text features
+            features_protein_all: [B * num_gpus, num_queries, D] - all protein features
+            features_text_all: [B * num_gpus, D] - all text features
+            return_sim: whether to return similarity matrices
+        """
+        bs = features_protein.size(0)
+        device = features_protein.device
+
+        # Protein-to-text similarity
+        # shape: [B, 1, num_queries, D] @ [B * num_gpus, D, 1] -> [B, B * num_gpus, num_queries]
+        sim_p2t = (features_protein.unsqueeze(1) @ features_text_all.unsqueeze(-1)).squeeze()
+        sim_p2t, _ = sim_p2t.max(-1)  # Take max over query tokens: [B, B * num_gpus]
+
+        logits_per_protein = sim_p2t / self.temperature
+    
+        # Text-to-protein similarity  
+        # shape: [B, 1, 1, D] @ [B*num_gpus, D, num_queries] -> [B, B*num_gpus, num_queries]
+        sim_t2p = (features_text.unsqueeze(1).unsqueeze(1) @ features_protein_all.permute(0, 2, 1)).squeeze()
+        sim_t2p, _ = sim_t2p.max(-1)  # Take max over query tokens: [B, B * num_gpus]
+        logits_per_text = sim_t2p / self.temperature
+
+        # Create labels for current rank
+        if dist.is_available() and dist.is_initialized():
+            rank = dist.get_rank()
+            labels = torch.linspace(rank * bs, rank * bs + bs - 1, bs, dtype=int).to(device)
+        else:
+            labels = torch.arange(bs, dtype=torch.long, device=device)
+        
+        # Compute contrastive losses
+        loss_protein = F.cross_entropy(logits_per_protein, labels)
+        loss_text = F.cross_entropy(logits_per_text, labels)
+        loss = (loss_protein + loss_text) / 2
+
+        if return_sim:
+            return logits_per_protein, logits_per_text, loss
+        else:
+            return loss
+
+    def compute_ptm_loss(self, protein_embeds, protein_mask, text_ids, text_mask, 
+                    query_tokens, tokenizer, qformer, sim_p2t, sim_t2p):
+        """
+        Compute protein-text matching loss.
+        修改以匹配标准 BertModel 的 API
+        """
+        batch_size = protein_embeds.size(0)
+        device = protein_embeds.device
+        
+        # Get world features for negative sampling
+        protein_embeds_world = pl_concat_all_gather(protein_embeds)
+        protein_mask_world = pl_concat_all_gather(protein_mask)
+        text_ids_world = pl_concat_all_gather(text_ids)
+        text_mask_world = pl_concat_all_gather(text_mask)
+        
+        with torch.no_grad():
+            if dist.is_available() and dist.is_initialized():
+                rank = dist.get_rank()
+            else:
+                rank = 0
+                
+            # Compute weights for negative sampling
+            weights_t2p = F.softmax(sim_t2p, dim=1) + 1e-4
+            weights_t2p[:, rank * batch_size : rank * batch_size + batch_size].fill_diagonal_(0)
+            
+            weights_p2t = F.softmax(sim_p2t, dim=1) + 1e-4
+            weights_p2t[:, rank * batch_size : rank * batch_size + batch_size].fill_diagonal_(0)
+
+        # Select negative proteins for each text
+        protein_embeds_neg = []
+        protein_mask_neg = []
+        for b in range(batch_size):
+            neg_idx = torch.multinomial(weights_t2p[b], 1).item()
+            protein_embeds_neg.append(protein_embeds_world[neg_idx])
+            protein_mask_neg.append(protein_mask_world[neg_idx])
+        
+        protein_embeds_neg = torch.stack(protein_embeds_neg, dim=0)
+        protein_mask_neg = torch.stack(protein_mask_neg, dim=0)
+
+        # Select negative texts for each protein
+        text_ids_neg = []
+        text_mask_neg = []
+        for b in range(batch_size):
+            neg_idx = torch.multinomial(weights_p2t[b], 1).item()
+            text_ids_neg.append(text_ids_world[neg_idx])
+            text_mask_neg.append(text_mask_world[neg_idx])
+
+        text_ids_neg = torch.stack(text_ids_neg, dim=0)
+        text_mask_neg = torch.stack(text_mask_neg, dim=0)
+
+        # Prepare inputs for PTM
+        text_ids_all = torch.cat([text_ids, text_ids, text_ids_neg], dim=0)  # pos, pos, neg
+        text_mask_all = torch.cat([text_mask, text_mask, text_mask_neg], dim=0)
+
+        # 获取 text embeddings
+        text_embeds_all = qformer.embeddings.word_embeddings(text_ids_all)
+        
+        # Expand query tokens for all samples
+        query_tokens_ptm = query_tokens.expand(text_ids_all.shape[0], -1, -1)
+        query_mask_ptm = torch.ones(query_tokens_ptm.size()[:-1], dtype=torch.long, device=device)
+        
+        # 方法1：只使用 query tokens 和 text，不直接编码 protein
+        # 这更符合你当前的 QFormer 架构
+        inputs_embeds = torch.cat([query_tokens_ptm, text_embeds_all], dim=1)
+        attention_mask_all = torch.cat([query_mask_ptm, text_mask_all], dim=1)
+        
+        # 确保序列长度不超过限制
+        max_length = qformer.config.max_position_embeddings
+        if attention_mask_all.size(1) > max_length:
+            # 截断 text 部分
+            excess = attention_mask_all.size(1) - max_length
+            if excess > 0:
+                inputs_embeds = inputs_embeds[:, :-excess, :]
+                attention_mask_all = attention_mask_all[:, :-excess]
+        
+        # Forward through QFormer - 使用标准 BERT API
+        output_ptm = qformer(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask_all,
+            return_dict=True,
+        )
+
+        # Extract query embeddings
+        pl_embeddings = output_ptm.last_hidden_state[:, :query_tokens_ptm.size(1), :]
+        
+        # 确保 ptm_head 存在并且维度正确
+        if not hasattr(self, 'ptm_head'):
+            hidden_size = pl_embeddings.size(-1)
+            self.ptm_head = nn.Linear(hidden_size, 2).to(device)
+        
+        pl_output = self.ptm_head(pl_embeddings)
+        logits = pl_output.mean(dim=1)  # [batch_size * 3, 2]
+
+        # Create labels: positive pairs get label 1, negative pairs get label 0
+        ptm_labels = torch.cat([
+            torch.ones(batch_size, dtype=torch.long),   # text-protein positive
+            torch.zeros(batch_size, dtype=torch.long),  # text-protein_neg negative  
+            torch.zeros(batch_size, dtype=torch.long)   # text_neg-protein negative
+        ], dim=0).to(device)
+        
+        loss_ptm = F.cross_entropy(logits, ptm_labels)
+        return loss_ptm
+
+
+class ContrastiveTrainer(Trainer):
+    """
+    Enhanced trainer for contrastive learning between proteins and text.
+    """
+    
+    def __init__(
+        self,
+        model: ProteinLLMModel,
+        args: ContrastiveTrainingArguments,
+        train_dataset: Optional[Dataset] = None,
+        eval_dataset: Optional[Dataset] = None,
+        data_collator: Optional[callable] = None,
+        **kwargs
+    ):
+        self.contrastive_loss = EnhancedContrastiveLoss(
+            temperature=args.temperature, 
+            enable_ptm=args.enable_ptm
+        )
+        self.freeze_protein_model = args.freeze_protein_model
+        self.freeze_text_model = args.freeze_text_model
+        self.protein_weight = args.protein_weight
+        self.text_weight = args.text_weight
+        self.enable_ptm = args.enable_ptm
+        self.ptm_weight = args.ptm_weight
+        
+        # Freeze models if specified
+        if self.freeze_protein_model:
+            for param in model.protein_model.parameters():
+                param.requires_grad = False
+                
+        if self.freeze_text_model:
+            for param in model.text_model.parameters():
+                param.requires_grad = False
+        
+        # Ensure projection layers are trainable
+        for param in model.protein_projection.parameters():
+            param.requires_grad = True
+            
+        super().__init__(
+            model=model,
+            args=args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            data_collator=data_collator,
+            **kwargs
+        )
+
+    def compute_loss(self, model, inputs, return_outputs=False,**kwargs):
+        """
+        Compute enhanced contrastive loss with optional PTM.
+        """
+# 检查模型是否被DataParallel包装
+        if hasattr(model, 'module'):
+            # 如果是DataParallel，使用.module访问原始模型
+            model = model.module
+        else:
+            # 否则直接使用模型
+            model = model
+
+        num_items_in_batch = kwargs.get('num_items_in_batch', None)
+
+        protein_sequences = inputs["protein_sequences"]
+        text_sequences = inputs["text_sequences"]
+        
+        # Get device from model
+        device = next(model.parameters()).device
+
+        # Get protein embeddings (before projection)
+        protein_embeds = model.get_protein_embeddings(protein_sequences)  # [B, seq_len, hidden]
+        
+        # Get protein features through projection (query tokens)
+        protein_features = model.get_protein_features(protein_sequences)  # [B, num_queries, embed_dim]
+        
+        # Get text features
+        text_features = model.get_text_features(text_sequences)  # [B, embed_dim]
+        
+        # Normalize features
+        protein_features = F.normalize(protein_features, p=2, dim=-1)
+        text_features = F.normalize(text_features, p=2, dim=-1)
+        
+        # Gather features from all processes for global contrastive learning
+        protein_features_all = pl_concat_all_gather(protein_features)
+        text_features_all = pl_concat_all_gather(text_features)
+        
+        # Compute contrastive loss
+        sim_p2t, sim_t2p, loss_contrastive = self.contrastive_loss.contrast_global(
+            protein_features, text_features, protein_features_all, text_features_all, return_sim=True
+        )
+        
+        total_loss = loss_contrastive
+        
+        # Compute PTM loss if enabled
+        loss_ptm = 0
+        if self.enable_ptm:
+            # Tokenize text for PTM
+            text_tokenized = model.text_tokenizer(
+                text_sequences, 
+                padding=True, 
+                truncation=True, 
+                return_tensors="pt",
+                max_length=self.args.max_length_text
+            ).to(model.device)
+            
+            # Get protein attention mask
+            protein_tokenized = model.protein_tokenizer(
+                protein_sequences,
+                padding=True,
+                truncation=True,
+                return_tensors="pt",
+                max_length=self.args.max_length_protein
+            ).to(model.device)
+            
+            loss_ptm = self.contrastive_loss.compute_ptm_loss(
+                protein_embeds=protein_embeds,
+                protein_mask=protein_tokenized.attention_mask,
+                text_ids=text_tokenized.input_ids,
+                text_mask=text_tokenized.attention_mask,
+                query_tokens=model.protein_projection.query_tokens,
+                tokenizer=model.text_tokenizer,
+                qformer=model.protein_projection.qformer,
+                sim_p2t=sim_p2t,
+                sim_t2p=sim_t2p
+            )
+            
+            total_loss = total_loss + self.ptm_weight * loss_ptm
+        
+        # Log detailed metrics
+        with torch.no_grad():
+            # Compute similarity statistics
+            similarity_matrix = torch.matmul(protein_features.flatten(0, 1), text_features.T)
+            positive_similarities = torch.diag(similarity_matrix[:protein_features.size(0)])
+            negative_similarities = similarity_matrix[~torch.eye(similarity_matrix.size(0), dtype=bool)]
+            
+            log_dict = {
+                "contrastive_loss": loss_contrastive.item(),
+                "total_loss": total_loss.item(),
+                "positive_similarity_mean": positive_similarities.mean().item(),
+                "negative_similarity_mean": negative_similarities.mean().item(),
+                "positive_similarity_std": positive_similarities.std().item(),
+                "similarity_gap": (positive_similarities.mean() - negative_similarities.mean()).item(),
+            }
+            
+            if self.enable_ptm:
+                log_dict["ptm_loss"] = loss_ptm.item() if isinstance(loss_ptm, torch.Tensor) else loss_ptm
+            
+            self.log(log_dict)
+        
+        if return_outputs:
+            outputs = {
+                "protein_features": protein_features,
+                "text_features": text_features,
+                "similarity_matrix_p2t": sim_p2t,
+                "similarity_matrix_t2p": sim_t2p,
+                "loss_contrastive": loss_contrastive,
+                "loss_ptm": loss_ptm,
+            }
+            return (total_loss, outputs)
+        
+        return total_loss
+
+    def evaluation_loop(self, dataloader, description, prediction_loss_only=None, ignore_keys=None, metric_key_prefix="eval"):
+        """
+        Enhanced evaluation loop for contrastive learning with PTM.
+        """
+        model = self._wrap_model(self.model, training=False, dataloader=dataloader)
+        model.eval()
+        
+        total_loss = 0.0
+        total_contrastive_loss = 0.0
+        total_ptm_loss = 0.0
+        total_samples = 0
+        all_protein_features = []
+        all_text_features = []
+        
+        for step, inputs in enumerate(dataloader):
+            with torch.no_grad():
+                loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
+                
+                total_loss += loss.item()
+                total_contrastive_loss += outputs["loss_contrastive"].item()
+                if self.enable_ptm:
+                    ptm_loss_val = outputs["loss_ptm"].item() if isinstance(outputs["loss_ptm"], torch.Tensor) else outputs["loss_ptm"]
+                    total_ptm_loss += ptm_loss_val
+                
+                total_samples += len(inputs["protein_sequences"])
+                
+                # Collect features for retrieval metrics
+                all_protein_features.append(outputs["protein_features"].cpu())
+                all_text_features.append(outputs["text_features"].cpu())
+        
+        # Compute average losses
+        avg_loss = total_loss / len(dataloader)
+        avg_contrastive_loss = total_contrastive_loss / len(dataloader)
+        avg_ptm_loss = total_ptm_loss / len(dataloader) if self.enable_ptm else 0
+        
+        # Concatenate all features for retrieval metrics
+        all_protein_features = torch.cat(all_protein_features, dim=0)
+        all_text_features = torch.cat(all_text_features, dim=0)
+        
+        # Compute retrieval metrics
+        retrieval_metrics = self.compute_retrieval_metrics(all_protein_features, all_text_features)
+        
+        metrics = {
+            f"{metric_key_prefix}_loss": avg_loss,
+            f"{metric_key_prefix}_contrastive_loss": avg_contrastive_loss,
+            **{f"{metric_key_prefix}_{k}": v for k, v in retrieval_metrics.items()}
+        }
+        
+        if self.enable_ptm:
+            metrics[f"{metric_key_prefix}_ptm_loss"] = avg_ptm_loss
+        
+        return metrics
+
+    def compute_retrieval_metrics(self, protein_features: torch.Tensor, text_features: torch.Tensor) -> Dict[str, float]:
+        """
+        Compute retrieval metrics for multi-query protein features.
+        """
+        # Handle multi-query protein features by taking mean or max
+        if protein_features.dim() == 3:  # [batch, num_queries, embed_dim]
+            protein_features_pooled = protein_features.mean(dim=1)  # Pool query tokens
+        else:
+            protein_features_pooled = protein_features
+        
+        # Normalize features
+        protein_features_pooled = F.normalize(protein_features_pooled, dim=-1)
+        text_features = F.normalize(text_features, dim=-1)
+        
+        # Compute similarity matrix
+        similarity_matrix = torch.matmul(protein_features_pooled, text_features.T)
+        
+        # Protein-to-text retrieval
+        p2t_ranks = []
+        for i in range(similarity_matrix.size(0)):
+            similarities = similarity_matrix[i]
+            rank = (similarities >= similarities[i]).sum().item()
+            p2t_ranks.append(rank)
+        
+        # Text-to-protein retrieval  
+        t2p_ranks = []
+        for i in range(similarity_matrix.size(1)):
+            similarities = similarity_matrix[:, i]
+            rank = (similarities >= similarities[i]).sum().item()
+            t2p_ranks.append(rank)
+        
+        # Compute Recall@K
+        metrics = {}
+        for k in [1, 5, 10]:
+            p2t_recall_k = sum(1 for rank in p2t_ranks if rank <= k) / len(p2t_ranks)
+            t2p_recall_k = sum(1 for rank in t2p_ranks if rank <= k) / len(t2p_ranks)
+            
+            metrics[f"p2t_recall_at_{k}"] = p2t_recall_k
+            metrics[f"t2p_recall_at_{k}"] = t2p_recall_k
+            metrics[f"avg_recall_at_{k}"] = (p2t_recall_k + t2p_recall_k) / 2
+        
+        # Mean rank
+        metrics["p2t_mean_rank"] = sum(p2t_ranks) / len(p2t_ranks)
+        metrics["t2p_mean_rank"] = sum(t2p_ranks) / len(t2p_ranks)
+        
+        return metrics
+
+
+def protein_text_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, List[str]]:
+    """
+    Collate function for protein-text contrastive learning.
+    """
+    protein_sequences = [item["protein_sequence"] for item in batch]
+    text_sequences = [item["text_description"] for item in batch]
+    
+    return {
+        "protein_sequences": protein_sequences,
+        "text_sequences": text_sequences,
+    }
+
+
+# Example usage function remains the same but with enhanced arguments
+# def train_contrastive_model(
+#     model: ProteinLLMModel,
+#     train_dataset: Dataset,
+#     eval_dataset: Optional[Dataset] = None,
+#     output_dir: str = "./contrastive_outputs",
+#     num_epochs: int = 10,
+#     batch_size: int = 32,
+#     learning_rate: float = 1e-4,
+#     temperature: float = 0.07,
+#     enable_ptm: bool = True,
+#     ptm_weight: float = 1.0,
+#     **kwargs
+# ):
+#     """
+#     Train the model with enhanced contrastive learning.
+#     """
+#     training_args = ContrastiveTrainingArguments(
+#         output_dir=output_dir,
+#         num_train_epochs=num_epochs,
+#         per_device_train_batch_size=batch_size,
+#         per_device_eval_batch_size=batch_size,
+#         learning_rate=learning_rate,
+#         temperature=temperature,
+#         enable_ptm=enable_ptm,
+#         ptm_weight=ptm_weight,
+#         logging_steps=10,
+#         evaluation_strategy="steps" if eval_dataset else "no",
+#         eval_steps=100 if eval_dataset else None,
+#         save_steps=500,
+#         save_total_limit=3,
+#         load_best_model_at_end=True if eval_dataset else False,
+#         metric_for_best_model="eval_avg_recall_at_1" if eval_dataset else None,
+#         greater_is_better=True,
+#         report_to=["wandb"] if wandb.run else [],
+#         **kwargs
+#     )
+    
+#     trainer = ContrastiveTrainer(
+#         model=model,
+#         args=training_args,
+#         train_dataset=train_dataset,
+#         eval_dataset=eval_dataset,
+#         data_collator=protein_text_collate_fn,
+#     )
+    
+#     # Train the model
+#     trainer.train()
+    
+#     # Save the final model
+#     trainer.save_model()
+    
+#     return trainer
+
+
+
+
+# def compute_ptm_loss(self, protein_embeds, protein_mask, text_ids, text_mask, 
+#                         query_tokens, tokenizer, qformer, sim_p2t, sim_t2p):
+#         """
+#         Compute protein-text matching loss.
+#         """
+#         batch_size = protein_embeds.size(0)
+#         device = protein_embeds.device
+        
+#         # Get world features for negative sampling
+#         protein_embeds_world = pl_concat_all_gather(protein_embeds)
+#         protein_mask_world = pl_concat_all_gather(protein_mask)
+#         text_ids_world = pl_concat_all_gather(text_ids)
+#         text_mask_world = pl_concat_all_gather(text_mask)
+        
+#         with torch.no_grad():
+#             if dist.is_available() and dist.is_initialized():
+#                 rank = dist.get_rank()
+#             else:
+#                 rank = 0
+                
+#             # Compute weights for negative sampling
+#             weights_t2p = F.softmax(sim_t2p, dim=1) + 1e-4
+#             weights_t2p[:, rank * batch_size : rank * batch_size + batch_size].fill_diagonal_(0)
+            
+#             weights_p2t = F.softmax(sim_p2t, dim=1) + 1e-4
+#             weights_p2t[:, rank * batch_size : rank * batch_size + batch_size].fill_diagonal_(0)
+
+#         # Select negative proteins for each text
+#         protein_embeds_neg = []
+#         protein_mask_neg = []
+#         for b in range(batch_size):
+#             neg_idx = torch.multinomial(weights_t2p[b], 1).item()
+#             protein_embeds_neg.append(protein_embeds_world[neg_idx])
+#             protein_mask_neg.append(protein_mask_world[neg_idx])
+        
+#         protein_embeds_neg = torch.stack(protein_embeds_neg, dim=0)
+#         protein_mask_neg = torch.stack(protein_mask_neg, dim=0)
+
+#         # Select negative texts for each protein
+#         text_ids_neg = []
+#         text_mask_neg = []
+#         for b in range(batch_size):
+#             neg_idx = torch.multinomial(weights_p2t[b], 1).item()
+#             text_ids_neg.append(text_ids_world[neg_idx])
+#             text_mask_neg.append(text_mask_world[neg_idx])
+
+#         text_ids_neg = torch.stack(text_ids_neg, dim=0)
+#         text_mask_neg = torch.stack(text_mask_neg, dim=0)
+
+#         # Prepare inputs for PTM
+#         text_ids_all = torch.cat([text_ids, text_ids, text_ids_neg], dim=0)  # pos, pos, neg
+#         text_mask_all = torch.cat([text_mask, text_mask, text_mask_neg], dim=0)
+
+#         query_tokens_ptm = query_tokens.expand(text_ids_all.shape[0], -1, -1)
+#         query_mask_ptm = torch.ones(query_tokens_ptm.size()[:-1], dtype=torch.long, device=device)
+#         attention_mask_all = torch.cat([query_mask_ptm, text_mask_all], dim=1)
+
+#         protein_embeds_all = torch.cat([protein_embeds, protein_embeds_neg, protein_embeds], dim=0)  # pos, neg, pos
+#         protein_mask_all = torch.cat([protein_mask, protein_mask_neg, protein_mask], dim=0)
+        
+#         # Combine embeddings
+#         inputs_embeds = torch.cat([query_tokens_ptm, 
+#                                  qformer.embeddings.word_embeddings(text_ids_all)], dim=1)
+        
+#         # Create position ids
+#         position_ids = torch.arange(
+#             attention_mask_all.size(1), dtype=torch.long, device=device
+#         ).unsqueeze(0).expand(text_ids_all.size(0), -1)
+        
+#         # Forward through QFormer for PTM - using BERT's forward method directly
+#         output_ptm = qformer(
+#             inputs_embeds=inputs_embeds,
+#             attention_mask=attention_mask_all,
+#             position_ids=position_ids,
+#             encoder_hidden_states=protein_embeds_all,
+#             encoder_attention_mask=protein_mask_all,
+#             return_dict=True,
+#         )
+
+#         pl_embeddings = output_ptm.last_hidden_state[:, :query_tokens_ptm.size(1), :]
+#         pl_output = self.ptm_head(pl_embeddings)
+#         logits = pl_output.mean(dim=1)
+
+#         ptm_labels = torch.cat([
+#             torch.ones(batch_size, dtype=torch.long), 
+#             torch.zeros(2 * batch_size, dtype=torch.long)
+#         ], dim=0).to(device)
+        
+#         loss_ptm = F.cross_entropy(logits, ptm_labels)
+#         return loss_ptm

BioReason_new/bioreason/trainer/grpo_config.py

@@ -0,0 +1,338 @@
+# 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 ProteinLLMGRPOConfig(TrainingArguments):
+    r"""
+    Configuration class for the [`ProteinLLMGRPOTrainer`].
+
+    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 [`ProteinLLMGRPOTrainer`] 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 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 `ProteinLLMGRPOTrainer` 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 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."
+            )
+        },
+    )
+
+
+    # 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."},
+    )

BioReason_new/bioreason/trainer/grpo_trainer.py

@@ -0,0 +1,719 @@
+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 (
+    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.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 accelerate.utils import is_peft_model, set_seed, gather_object
+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.protein_modules.protein_base_module import ProteinBaseModule
+from bioreason.trainer.protein_grpo_config import ProteinLLMGRPOConfig
+
+# 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.
+    """
+
+    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 ProteinLLMGRPOTrainer(Trainer):
+    """
+    Trainer for the Group Relative Policy Optimization (GRPO) method adapted for Protein-LLM models.
+    """
+
+    def __init__(
+        self,
+        model: Union[str, PreTrainedModel],
+        reward_funcs: Union[RewardFunc, list[RewardFunc]],
+        args: ProteinLLMGRPOConfig = None,
+        protein_module: ProteinBaseModule = 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_protein_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.protein_module = protein_module
+
+        # Models
+        model_init_kwargs = args.model_init_kwargs or {}
+        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
+        elif isinstance(torch_dtype, str):
+            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}."
+            )
+
+        model_init_kwargs["use_cache"] = (
+            False if args.gradient_checkpointing else model_init_kwargs.get("use_cache")
+        )
+
+        # LoRA
+        self.protein_modules_keywords = self.protein_module.get_proteinllm_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 protein 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:
+                    if "embed_tokens" in m:
+                        lora_module_names.remove(m)
+                return list(lora_module_names)
+            target_modules = find_all_linear_names(model, self.protein_modules_keywords)
+            peft_config.target_modules = target_modules
+            model = prepare_model_for_kbit_training(model)
+            model = get_peft_model(model, peft_config)
+
+        # Freeze protein modules
+        if freeze_protein_modules:
+            print("Freezing protein modules...")
+            for p in model.protein_model.parameters():
+                p.requires_grad = False
+        
+        # Make projection layer trainable
+        for p in model.protein_projection.parameters():
+            p.requires_grad = True
+
+        # Compute the number of trainable parameters
+        trainable_params = [p for p in model.parameters() if p.requires_grad]
+        total_params = sum(p.numel() for p in trainable_params)
+        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:
+            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):
+            self.ref_model = None
+        else:
+            self.ref_model = create_reference_model(model)
+
+        # Processing class
+        if processing_class is None:
+            processing_cls = self.protein_module.get_processing_class()
+            processing_class = processing_cls(
+                tokenizer=model.text_tokenizer, 
+                protein_tokenizer=model.protein_tokenizer
+            )
+            
+            for component, processing_keyword in self.protein_module.get_custom_processing_keywords():
+                if processing_keyword in kwargs:
+                    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.protein_module.post_model_init(model, processing_class)
+        if self.ref_model is not None:
+            self.protein_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
+                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):
+            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
+        self.num_generations = args.num_generations
+        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.protein_module, "get_eos_token_id"):
+            self.generation_config.eos_token_id = self.protein_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
+        self._step = 0
+        self._buffered_inputs = [None] * args.gradient_accumulation_steps
+
+        # Suppress warnings
+        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 batch size compatibility
+        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}."
+                )
+
+        # Set seed for reproducibility
+        set_seed(args.seed, device_specific=True)
+
+        # Gradient accumulation setup
+        self.model_accepts_loss_kwargs = False
+
+        if self.ref_model is not None:
+            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: ProteinLLMGRPOConfig) -> PreTrainedModel:
+        """Enables gradient checkpointing for the model."""
+        model.config.use_cache = False
+
+        if is_peft_model(model):
+            model.base_model.gradient_checkpointing_enable()
+        else:
+            if getattr(model, "language_model", None) is not None:
+                model.language_model.config.use_cache = False
+                model.protein_model.gradient_checkpointing = True
+                model.language_model._set_gradient_checkpointing()
+                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._signature_columns is None:
+            self._signature_columns = ["prompt"]
+
+    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
+        logits = logits[:, :-1, :]
+        input_ids = input_ids[:, 1:]
+        
+        # Compute the log probabilities for the input tokens
+        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):
+        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.protein_module.prepare_prompt(self.processing_class, inputs)
+        
+        # Handle protein sequences
+        batch_protein_sequences = []
+        print("_generate_and_score_completions (GRPO):")
+        for x in inputs:
+            if 'protein_sequences' in x:
+                proteins = self._get_key_from_inputs(x, "protein_sequences")
+            else:
+                proteins = []
+            batch_protein_sequences.append(proteins)
+
+        prompt_inputs = self.protein_module.prepare_model_inputs(
+            self.processing_class,
+            model,
+            prompts_text,
+            batch_protein_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"]
+
+        # 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.protein_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.protein_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:
+                completion_ids = generate_returned_result
+                prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1)
+
+        # 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)
+
+        # Get the multimodal inputs
+        multimodal_keywords = self.protein_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():
+            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
+        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]
+            else:
+                reward_kwargs = {key: [] for key in inputs[0].keys() if key not in ["prompt", "completion"]}
+                for key in reward_kwargs:
+                    for example in inputs:
+                        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
+        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())
+
+        if (
+            self.log_completions
+            and self.state.global_step % self.args.logging_steps == 0
+            and "wandb" in self.args.report_to
+        ):
+            timestamp = time.time()
+            num_items = len(gather_object(prompts_text))
+
+            table = {
+                "step": [f"{self.state.global_step}_{timestamp}"] * num_items,
+                "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 ProteinLLMGRPOTrainer does not support returning outputs")
+    
+        # Check if we need to generate new completions or use buffered ones
+        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
+        
+        # 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)
+        
+        # Get the current policy's log probabilities
+        per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, **multimodal_inputs)
+        per_token_logps = per_token_logps[:, prompt_ids.size(1) - 1:]
+
+        # Get the advantages from inputs
+        advantages = inputs["advantages"]
+        
+        # When using num_iterations == 1, old_per_token_logps == per_token_logps
+        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)
+        
+        # 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
+        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()}
+        logs = {**logs, **metrics}
+        if version.parse(transformers.__version__) >= version.parse("4.47.0.dev0"):
+            super().log(logs, start_time)
+        else:
+            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`.
+        """
+        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_new/bioreason/utils/protein_utils.py

@@ -0,0 +1,229 @@
+# 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
+
+# ProteinInput = Union[
+#     str, list[int], np.ndarray, "torch.Tensor", list[str], list[list[int]], list[np.ndarray], list["torch.Tensor"]
+# ]  # noqa
+
+
+# def clean_protein_sequence(sequence: str) -> str:
+#     """
+#     Clean protein sequence by removing invalid characters and normalizing.
+    
+#     Args:
+#         sequence: Raw protein sequence string
+        
+#     Returns:
+#         Cleaned protein sequence
+#     """
+#     # Remove whitespace and convert to uppercase
+#     sequence = sequence.replace(" ", "").replace("\n", "").upper()
+    
+#     # Keep only valid amino acid characters
+#     valid_aa = set("ACDEFGHIKLMNPQRSTVWY")
+#     cleaned_sequence = "".join(char for char in sequence if char in valid_aa)
+    
+#     return cleaned_sequence
+
+
+# def truncate_protein_sequence(sequence: str, max_length: int = 1024) -> str:
+#     """
+#     Truncate protein sequence to maximum length.
+    
+#     Args:
+#         sequence: Protein sequence string
+#         max_length: Maximum allowed length
+        
+#     Returns:
+#         Truncated protein sequence
+#     """
+#     if len(sequence) <= max_length:
+#         return sequence
+    
+#     # Truncate from both ends to keep the middle part (often most important)
+#     if max_length >= 100:
+#         start_keep = max_length // 3
+#         end_keep = max_length - start_keep
+#         return sequence[:start_keep] + sequence[-end_keep:]
+#     else:
+#         # If very short max_length, just truncate from end
+#         return sequence[:max_length]
+
+
+# def validate_protein_sequence(sequence: str) -> bool:
+#     """
+#     Validate if a sequence contains only valid amino acid characters.
+    
+#     Args:
+#         sequence: Protein sequence string
+        
+#     Returns:
+#         True if valid, False otherwise
+#     """
+#     valid_aa = set("ACDEFGHIKLMNPQRSTVWY")
+#     return all(char in valid_aa for char in sequence.upper())
+
+
+# def get_sequence_stats(sequence: str) -> dict:
+#     """
+#     Get basic statistics about a protein sequence.
+    
+#     Args:
+#         sequence: Protein sequence string
+        
+#     Returns:
+#         Dictionary with sequence statistics
+#     """
+#     sequence = sequence.upper()
+#     length = len(sequence)
+    
+#     if length == 0:
+#         return {"length": 0, "composition": {}, "molecular_weight": 0.0}
+    
+#     # Amino acid composition
+#     composition = {}
+#     for aa in "ACDEFGHIKLMNPQRSTVWY":
+#         count = sequence.count(aa)
+#         composition[aa] = {
+#             "count": count,
+#             "frequency": count / length if length > 0 else 0
+#         }
+    
+#     # Approximate molecular weight (Da)
+#     aa_weights = {
+#         'A': 89.1, 'C': 121.0, 'D': 133.1, 'E': 147.1, 'F': 165.2,
+#         'G': 75.1, 'H': 155.2, 'I': 131.2, 'K': 146.2, 'L': 131.2,
+#         'M': 149.2, 'N': 132.1, 'P': 115.1, 'Q': 146.2, 'R': 174.2,
+#         'S': 105.1, 'T': 119.1, 'V': 117.1, 'W': 204.2, 'Y': 181.2
+#     }
+    
+#     molecular_weight = sum(aa_weights.get(aa, 0) for aa in sequence)
+#     # Subtract water molecules for peptide bonds
+#     molecular_weight -= (length - 1) * 18.015 if length > 1 else 0
+    
+#     return {
+#         "length": length,
+#         "composition": composition,
+#         "molecular_weight": molecular_weight
+#     }
+
+
+# def format_protein_for_display(sequence: str, line_length: int = 80) -> str:
+#     """
+#     Format protein sequence for display with line breaks.
+    
+#     Args:
+#         sequence: Protein sequence string
+#         line_length: Number of characters per line
+        
+#     Returns:
+#         Formatted sequence string
+#     """
+#     if not sequence:
+#         return ""
+    
+#     lines = []
+#     for i in range(0, len(sequence), line_length):
+#         line = sequence[i:i + line_length]
+#         # Add position numbers
+#         pos_start = i + 1
+#         pos_end = min(i + line_length, len(sequence))
+#         lines.append(f"{pos_start:>8} {line} {pos_end}")
+    
+#     return "\n".join(lines)
+
+
+# def compare_protein_sequences(seq1: str, seq2: str) -> dict:
+#     """
+#     Compare two protein sequences and return similarity metrics.
+    
+#     Args:
+#         seq1: First protein sequence
+#         seq2: Second protein sequence
+        
+#     Returns:
+#         Dictionary with comparison metrics
+#     """
+#     seq1 = seq1.upper().replace(" ", "")
+#     seq2 = seq2.upper().replace(" ", "")
+    
+#     if not seq1 or not seq2:
+#         return {"identity": 0.0, "similarity": 0.0, "gaps": 0}
+    
+#     # Simple identity calculation (without proper alignment)
+#     min_len = min(len(seq1), len(seq2))
+#     max_len = max(len(seq1), len(seq2))
+    
+#     identical = 0
+#     for i in range(min_len):
+#         if seq1[i] == seq2[i]:
+#             identical += 1
+    
+#     identity = identical / max_len if max_len > 0 else 0.0
+#     gaps = abs(len(seq1) - len(seq2))
+    
+#     # Simple similarity (identical positions / shorter sequence length)
+#     similarity = identical / min_len if min_len > 0 else 0.0
+    
+#     return {
+#         "identity": identity,
+#         "similarity": similarity, 
+#         "gaps": gaps,
+#         "identical_positions": identical,
+#         "seq1_length": len(seq1),
+#         "seq2_length": len(seq2)
+#     }
+
+
+# def extract_protein_domains(sequence: str, domain_patterns: dict = None) -> list:
+#     """
+#     Extract potential protein domains based on simple patterns.
+    
+#     Args:
+#         sequence: Protein sequence string
+#         domain_patterns: Dictionary of domain name to regex pattern
+        
+#     Returns:
+#         List of detected domains
+#     """
+#     import re
+    
+#     if domain_patterns is None:
+#         # Simple example patterns (in real use, you'd use proper domain databases)
+#         domain_patterns = {
+#             "signal_peptide": r"^M[A-Z]{10,30}[RK]",  # Very simple signal peptide pattern
+#             "transmembrane": r"[AILMFWYV]{15,25}",     # Hydrophobic stretch
+#             "nuclear_localization": r"[KR]{2,}[A-Z]{10,20}[KR]{2,}",  # Basic NLS pattern
+#         }
+    
+#     domains = []
+#     for domain_name, pattern in domain_patterns.items():
+#         matches = list(re.finditer(pattern, sequence))
+#         for match in matches:
+#             domains.append({
+#                 "domain": domain_name,
+#                 "start": match.start() + 1,  # 1-based indexing
+#                 "end": match.end(),
+#                 "sequence": match.group()
+#             })
+    
+#     return domains
+
+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
+
+ProteinInput = Union[
+    str, list[int], np.ndarray, "torch.Tensor", list[str], list[list[int]], list[np.ndarray], list["torch.Tensor"]
+]  # noqa

BioReason_new/readme.md

@@ -0,0 +1,8 @@
+# 1. 对比学习预训练
+python train_contrastive.py --use_wandb --freeze_protein_model --freeze_text_model
+
+# 2. 监督微调  
+python train_protein_qwen.py --model_type protein-llm --text_model_finetune True
+
+# 3. GRPO训练
+python protein_reason.py --sft_checkpoint ./checkpoints/best_model

BioReason_new/reason.py

@@ -0,0 +1,520 @@
+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
+
+# Import BLIP2 modules
+from model.blip2_stage2 import Blip2Stage2
+from blip2_dna_module import Blip2DNAModule
+from blip2_grpo_trainer import Blip2GRPOTrainer
+from bioreason.trainer import DNALLMGRPOConfig
+
+# 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)
+        
+        # For BLIP2, save the config from the LLM component
+        if hasattr(unwrapped_model, "blip2") and hasattr(unwrapped_model.blip2, "llm_model"):
+            if hasattr(unwrapped_model.blip2.llm_model, "config"):
+                unwrapped_model.blip2.llm_model.config.save_pretrained(checkpoint_folder)
+            elif hasattr(unwrapped_model.blip2.llm_model, "base_model") and hasattr(unwrapped_model.blip2.llm_model.base_model, "config"):
+                unwrapped_model.blip2.llm_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 extract_xml_answer(text: str) -> str:
+    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
+
+    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
+
+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
+
+# Reward functions
+def format_correct_reward_func(completions, **kwargs) -> list[float]:
+    """
+    奖励函数：检查格式是否正确
+    要求��包含 <think>...</think> 和 <answer>...</answer> 标签
+    """
+    responses = [completion[0]["content"] for completion in completions]
+    rewards = []
+    
+    for response in responses:
+        score = 0.0
+        
+        # 检查是否有think标签
+        if "<think>" in response and "</think>" in response:
+            score += 0.5
+            
+        # 检查是否有answer标签  
+        if "<answer>" in response and "</answer>" in response:
+            score += 0.5
+            
+        # 检查标签的顺序是否正确
+        think_start = response.find("<think>")
+        think_end = response.find("</think>")
+        answer_start = response.find("<answer>")
+        answer_end = response.find("</answer>")
+        
+        if (think_start != -1 and think_end != -1 and 
+            answer_start != -1 and answer_end != -1 and
+            think_start < think_end < answer_start < answer_end):
+            score += 0.5  # 格式完全正确的额外奖励
+            
+        rewards.append(score)
+    
+    return rewards
+
+def accuracy_reward_func(prompts, completions, answer, **kwargs) -> list[float]:
+    """
+    奖励函数：检查答案准确率
+    """
+    responses = [completion[0]['content'] for completion in completions]
+    rewards = []
+    
+    for i, response in enumerate(responses):
+        # 提取answer标签中的内容
+        answer_match = re.search(r"<answer>(.*?)</answer>", response, re.DOTALL)
+        if answer_match:
+            extracted_answer = answer_match.group(1).strip()
+        else:
+            extracted_answer = response.strip()
+        
+        # 获取正确答案
+        if isinstance(answer, list) and len(answer) > i:
+            correct_answer = str(answer[i]).strip()
+        elif isinstance(answer, list) and len(answer) > 0:
+            correct_answer = str(answer[0]).strip()
+        else:
+            correct_answer = str(answer).strip()
+        
+        # 计算准确率奖励
+        if correct_answer.lower() in extracted_answer.lower():
+            rewards.append(1.0)  # 完全匹配
+        elif any(word in extracted_answer.lower() for word in correct_answer.lower().split()):
+            rewards.append(0.5)  # 部分匹配
+        else:
+            rewards.append(0.0)  # 不匹配
+    
+    return rewards
+
+def repetition_penalty_reward_func(completions, **kwargs) -> list[float]:
+    """
+    奖励函数：检查重复率（越低越好）
+    计算文本中重复词汇的比例，重复率越低奖励越高
+    """
+    responses = [completion[0]["content"] for completion in completions]
+    rewards = []
+    
+    for response in responses:
+        # 提取answer部分的文本
+        answer_match = re.search(r"<answer>(.*?)</answer>", response, re.DOTALL)
+        if answer_match:
+            text_to_analyze = answer_match.group(1).strip()
+        else:
+            text_to_analyze = response.strip()
+        
+        # 分词并计算重复率
+        words = text_to_analyze.lower().split()
+        
+        if len(words) == 0:
+            rewards.append(0.0)
+            continue
+            
+        # 计算词汇重复率
+        unique_words = set(words)
+        repetition_rate = 1.0 - (len(unique_words) / len(words))
+        
+        # 计算句子重复率
+        sentences = [s.strip() for s in text_to_analyze.split('.') if s.strip()]
+        if len(sentences) > 1:
+            unique_sentences = set(sentences)
+            sentence_repetition_rate = 1.0 - (len(unique_sentences) / len(sentences))
+        else:
+            sentence_repetition_rate = 0.0
+        
+        # 综合重复率
+        overall_repetition = (repetition_rate + sentence_repetition_rate) / 2
+        
+        # 重复率越低，奖励越高
+        reward = max(0.0, 1.0 - overall_repetition * 2)  # 乘以2让惩罚更明显
+        rewards.append(reward)
+    
+    return rewards
+
+def combined_reward_func(prompts, completions, answer, 
+                        format_weight=0.3, accuracy_weight=0.5, repetition_weight=0.2,
+                        **kwargs) -> list[float]:
+    """
+    组合奖励函数：格式+准确率+重复率的加权组合
+    """
+    format_rewards = format_correct_reward_func(completions, **kwargs)
+    accuracy_rewards = accuracy_reward_func(prompts, completions, answer, **kwargs)
+    repetition_rewards = repetition_penalty_reward_func(completions, **kwargs)
+    
+    # 确保权重总和为1
+    total_weight = format_weight + accuracy_weight + repetition_weight
+    if total_weight != 1.0:
+        format_weight /= total_weight
+        accuracy_weight /= total_weight  
+        repetition_weight /= total_weight
+        print(f"Normalized weights - Format: {format_weight:.3f}, Accuracy: {accuracy_weight:.3f}, Repetition: {repetition_weight:.3f}")
+    
+    combined_rewards = []
+    for f_reward, a_reward, r_reward in zip(format_rewards, accuracy_rewards, repetition_rewards):
+        combined = (format_weight * f_reward + 
+                   accuracy_weight * a_reward + 
+                   repetition_weight * r_reward)
+        combined_rewards.append(combined)
+    
+    return combined_rewards
+
+# 保留一些原有的奖励函数作为备选
+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 xmlcount_reward_func(completions, **kwargs) -> list[float]:
+    contents = [completion[0]["content"] for completion in completions]
+    return [count_xml(c) for c in contents]
+
+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
+
+@dataclass
+class Blip2ModelConfig(ModelConfig):
+    # BLIP2 specific configuration
+    model_name_or_path: str = field(default="blip2-model", metadata={"help": "Model checkpoint for weights initialization."})
+    
+    # BLIP2 Architecture parameters
+    bert_name: str = field(default="/path/to/bert", metadata={"help": "BERT model for Q-former"})
+    num_query_token: int = field(default=32, metadata={"help": "Number of query tokens"})
+    cross_attention_freq: int = field(default=2, metadata={"help": "Cross attention frequency"})
+    plm_model: str = field(default="facebook/esm2_t30_150M_UR50D", metadata={"help": "Protein language model"})
+    plm_tune: str = field(default="freeze", metadata={"help": "PLM tuning strategy"})
+    llm_name: str = field(default="facebook/galactica-1.3b", metadata={"help": "Language model name"})
+    llm_tune: str = field(default="lora", metadata={"help": "LLM tuning strategy"})
+    qformer_tune: str = field(default="train", metadata={"help": "Q-former tuning strategy"})
+    peft_dir: str = field(default="", metadata={"help": "PEFT directory"})
+    
+    # LoRA parameters
+    lora_r: int = field(default=8, metadata={"help": "LoRA rank"})
+    lora_alpha: int = field(default=16, metadata={"help": "LoRA alpha"})
+    lora_dropout: float = field(default=0.1, metadata={"help": "LoRA dropout"})
+    
+    # Training parameters
+    enbale_gradient_checkpointing: bool = field(default=False, metadata={"help": "Enable gradient checkpointing"})
+    enable_flash: bool = field(default=False, metadata={"help": "Enable flash attention"})
+    
+    # Other parameters
+    cache_dir: str = field(default=None, metadata={"help": "Path to model cache directory."})
+    sft_checkpoint: str = field(default=None, metadata={"help": "Path to the checkpoint for SFT."})
+    freeze_dna_modules: bool = field(default=False, metadata={"help": "Freeze DNA/protein modules"})
+
+@dataclass
+class GRPOScriptArguments(ScriptArguments):
+    """
+    Script arguments for the GRPO training script with BLIP2.
+    """
+    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(
+        # 选项1：使用组合奖励函数（推荐）
+        default_factory=lambda: ["combined"],
+        
+        # 选项2：使用分离的三个奖励函数
+        # default_factory=lambda: ["format_correct", "accuracy", "repetition_penalty"],
+        
+        # 选项3：自定义组合
+        # default_factory=lambda: ["format_correct", "accuracy", "repetition_penalty", "xmlcount"],
+        
+        metadata={"help": "List of reward functions. Available: 'combined', 'format_correct', 'accuracy', 'repetition_penalty', 'xmlcount', 'strict_format', 'less_than_4'"},
+    )
+    
+    # 奖励函数权重配置
+    format_weight: float = field(
+        default=0.3,
+        metadata={"help": "Weight for format correctness reward (used in combined reward)"}
+    )
+    accuracy_weight: float = field(
+        default=0.5,
+        metadata={"help": "Weight for accuracy reward (used in combined reward)"}
+    )
+    repetition_weight: float = field(
+        default=0.2,
+        metadata={"help": "Weight for repetition penalty reward (used in combined reward)"}
+    )
+
+reward_funcs_registry = {
+    # 新的三合一奖励函数
+    "combined": combined_reward_func,          # 格式+准确率+重复率组合
+    
+    # 分离的奖励函数
+    "format_correct": format_correct_reward_func,      # 格式正确性
+    "accuracy": accuracy_reward_func,                  # 准确率
+    "repetition_penalty": repetition_penalty_reward_func,  # 重复率惩罚
+    
+    # 原有的奖励函数（保留作为备选）
+    "xmlcount": xmlcount_reward_func,
+    "strict_format": strict_format_reward_func,
+    "less_than_4": less_than_4_reward_func,
+}
+
+def get_vlm_module(model_name_or_path):
+    # Always use BLIP2 module for this implementation
+    return Blip2DNAModule
+
+def create_blip2_args_from_config(model_args):
+    """Create BLIP2 args from model config"""
+    # Convert model config to the format expected by BLIP2
+    blip2_args = {
+        'bert_name': model_args.bert_name,
+        'num_query_token': model_args.num_query_token,
+        'cross_attention_freq': model_args.cross_attention_freq,
+        'plm_model': model_args.plm_model,
+        'plm_tune': model_args.plm_tune,
+        'llm_name': model_args.llm_name,
+        'llm_tune': model_args.llm_tune,
+        'qformer_tune': model_args.qformer_tune,
+        'peft_dir': model_args.peft_dir,
+        'lora_r': model_args.lora_r,
+        'lora_alpha': model_args.lora_alpha,
+        'lora_dropout': model_args.lora_dropout,
+        'enbale_gradient_checkpointing': model_args.enbale_gradient_checkpointing,
+        'enable_flash': model_args.enable_flash,
+    }
+    return blip2_args
+
+def _prep_for_training(model, training_args):
+    """
+    Prepare BLIP2 model for training with LoRA.
+    """
+    # The BLIP2 model should handle its own LoRA setup
+    # This is mainly for any additional preparation needed
+    
+    target_modules = ["q_proj", "k_proj", "v_proj", "out_proj", "fc1", "fc2"]
+    
+    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",
+    )
+    
+    return lora_config
+
+def main(script_args, training_args, model_args):
+    print(training_args.output_dir)
+    torch.cuda.empty_cache()
+    torch.set_float32_matmul_precision("medium")
+
+    # Create BLIP2 model
+    blip2_args = create_blip2_args_from_config(model_args)
+    model = Blip2Stage2(blip2_args)
+
+    # Load checkpoint if specified
+    if model_args.sft_checkpoint is not None:
+        print(f"Loading SFT checkpoint from {model_args.sft_checkpoint}")
+        
+        if os.path.isdir(model_args.sft_checkpoint):
+            # Load Lightning checkpoint
+            checkpoint = torch.load(os.path.join(model_args.sft_checkpoint, "last.ckpt"), map_location='cpu')
+            model.load_state_dict(checkpoint['state_dict'], strict=False)
+            print("Loaded Lightning checkpoint")
+        else:
+            # Load PyTorch state dict
+            checkpoint = torch.load(model_args.sft_checkpoint, map_location='cpu')
+            
+            if "state_dict" in checkpoint:
+                state_dict = checkpoint["state_dict"]
+            else:
+                state_dict = checkpoint
+                
+            # Remove module prefix if present
+            state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+            
+            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")
+
+    # Get reward functions with weights
+    reward_funcs = []
+    for func_name in script_args.reward_funcs:
+        if func_name == "combined":
+            # 为组合奖励函数传递权重参数
+            def weighted_combined_reward(prompts, completions, answer, **kwargs):
+                return combined_reward_func(
+                    prompts, completions, answer,
+                    format_weight=script_args.format_weight,
+                    accuracy_weight=script_args.accuracy_weight,
+                    repetition_weight=script_args.repetition_weight,
+                    **kwargs
+                )
+            reward_funcs.append(weighted_combined_reward)
+        else:
+            reward_funcs.append(reward_funcs_registry[func_name])
+    
+    print("reward_funcs:", [func.__name__ if hasattr(func, '__name__') else 'weighted_combined_reward' for func in reward_funcs])
+    print(f"Reward weights - Format: {script_args.format_weight}, Accuracy: {script_args.accuracy_weight}, Repetition: {script_args.repetition_weight}")
+
+    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()
+
+    # Load dataset
+    dataset = get_kegg_questions()
+    print(dataset)
+
+    # Custom callback to handle saving with PyTorch's native mechanism
+    custom_save_callback = SaveWithPyTorchCallback()
+
+    # Initialize the BLIP2 GRPO trainer
+    trainer = Blip2GRPOTrainer(
+        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=getattr(model_args, 'attn_implementation', 'flash_attention_2'),
+        torch_dtype=getattr(model_args, 'torch_dtype', 'bfloat16'),
+        callbacks=[custom_save_callback],
+    )
+
+    # Set the trainer to save in PyTorch format instead of safetensors
+    training_args.save_safetensors = False
+
+    # Train the model
+    trainer.train()
+
+if __name__ == "__main__":
+    print(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES')}")
+    parser = TrlParser((GRPOScriptArguments, DNALLMGRPOConfig, Blip2ModelConfig))
+    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)

BioReason_new/run.sh

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

BioReason_new/run_contrast.sh

@@ -0,0 +1,31 @@
+echo "Starting contrastive pre-training..."
+export WANDB_BASE_URL=https://api.bandw.top
+
+# 指定要使用的 GPU 卡（例如使用 0,1,2,3 四张卡）
+export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+python train_contrastive.py \
+    --text_model_name "/oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged" \
+    --protein_model_name "/nas/shared/kilab/wangyujia/ProtT3/plm_model/esm2-150m" \
+    --qformer_model_name "/nas/shared/kilab/wangyujia/ProtT3/plm_model/microsoft" \
+    --num_query_tokens 8 \
+    --train_dataset "/nas/shared/kilab/wangyujia/ProtT3/data/SwissProtV3/train_set.jsonl" \
+    --valid_dataset "/nas/shared/kilab/wangyujia/ProtT3/data/SwissProtV3/valid_set.jsonl" \
+    --output_dir "./contrastive_outputs" \
+    --num_epochs 10 \
+    --batch_size 8 \
+    --learning_rate 1e-4 \
+    --temperature 0.07 \
+    --freeze_protein_model \
+    --freeze_text_model \
+    --enable_ptm \
+    --max_length_protein 1024 \
+    --max_length_text 512 \
+    --num_workers 8 \
+    --eval_dataset \
+    --use_wandb \
+    --wandb_project "protein-llm-contrastive" \
+    --logging_steps 100 \
+    --eval_steps 500 \
+    --save_steps 1000
+
+echo "Contrastive pre-training completed!"

BioReason_new/train_contrastive.py

@@ -0,0 +1,552 @@
+# import os
+# import time
+# from argparse import ArgumentParser
+# from functools import partial
+
+# import torch
+# import wandb
+# from datasets import load_dataset
+# from torch.utils.data import DataLoader
+
+# 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.models.protein_llm import ProteinLLMModel
+# from bioreason.models.contrast_trainer import (
+#     ContrastiveTrainer,
+#     ContrastiveTrainingArguments,
+#     protein_text_collate_fn,
+# )
+# from bioreason.dataset.protein import format_protein_contrastive
+
+
+# def main(args):
+#     """
+#     Main function for contrastive pre-training of Protein-LLM.
+    
+#     This script trains the QFormer projection layer to align protein and text representations
+#     using contrastive learning before fine-tuning on downstream tasks.
+#     """
+    
+#     # Set random seed
+#     pl.seed_everything(args.seed)
+#     torch.cuda.empty_cache()
+#     torch.set_float32_matmul_precision("medium")
+    
+#     # Initialize wandb
+#     if args.use_wandb:
+#         wandb.init(
+#             project=args.wandb_project,
+#             entity=args.wandb_entity,
+#             name=f"contrastive-{args.text_model_name.split('/')[-1]}-{time.strftime('%Y%m%d-%H%M%S')}",
+#             config=vars(args)
+#         )
+    
+#     print("Loading model...")
+#     # Load the Protein-LLM model
+#     model = ProteinLLMModel(
+#         text_model_name=args.text_model_name,
+#         protein_model_name=args.protein_model_name,
+#         qformer_model_name=args.qformer_model_name,
+#         cache_dir=args.cache_dir,
+#         max_length_protein=args.max_length_protein,
+#         max_length_text=args.max_length_text,
+#         text_model_finetune=False,  # Don't fine-tune during contrastive learning
+#         protein_model_finetune=False,  # Don't fine-tune during contrastive learning
+#         num_query_tokens=args.num_query_tokens,
+#     )
+    
+#     print("Loading datasets...")
+#     # Load datasets for contrastive learning
+#     train_dataset = load_dataset(args.dataset_name, split="train")
+#     eval_dataset = load_dataset(args.dataset_name, split="validation") if args.eval_dataset else None
+    
+#     # Format datasets for contrastive learning
+#     train_dataset = train_dataset.map(format_protein_contrastive)
+#     if eval_dataset:
+#         eval_dataset = eval_dataset.map(format_protein_contrastive)
+    
+#     # Filter out examples without protein sequences or descriptions
+#     train_dataset = train_dataset.filter(
+#         lambda x: x["protein_sequence"] and x["text_description"] 
+#         and len(x["protein_sequence"].strip()) > 0 
+#         and len(x["text_description"].strip()) > 0
+#     )
+    
+#     if eval_dataset:
+#         eval_dataset = eval_dataset.filter(
+#             lambda x: x["protein_sequence"] and x["text_description"] 
+#             and len(x["protein_sequence"].strip()) > 0 
+#             and len(x["text_description"].strip()) > 0
+#         )
+    
+#     print(f"Training dataset size: {len(train_dataset)}")
+#     if eval_dataset:
+#         print(f"Eval dataset size: {len(eval_dataset)}")
+    
+#     # Setup training arguments for contrastive learning
+#     training_args = ContrastiveTrainingArguments(
+#         output_dir=args.output_dir,
+#         num_train_epochs=args.num_epochs,
+#         per_device_train_batch_size=args.batch_size,
+#         per_device_eval_batch_size=args.batch_size,
+#         learning_rate=args.learning_rate,
+#         weight_decay=args.weight_decay,
+#         temperature=args.temperature,
+#         freeze_protein_model=args.freeze_protein_model,
+#         freeze_text_model=args.freeze_text_model,
+#         protein_weight=args.protein_weight,
+#         text_weight=args.text_weight,
+#         max_length_protein=args.max_length_protein,
+#         max_length_text=args.max_length_text,
+#         logging_steps=args.logging_steps,
+#         evaluation_strategy="steps" if eval_dataset else "no",
+#         eval_steps=args.eval_steps if eval_dataset else None,
+#         save_steps=args.save_steps,
+#         save_total_limit=args.save_total_limit,
+#         load_best_model_at_end=True if eval_dataset else False,
+#         metric_for_best_model="eval_avg_recall_at_1" if eval_dataset else None,
+#         greater_is_better=True,
+#         report_to=["wandb"] if args.use_wandb else [],
+#         warmup_steps=args.warmup_steps,
+#         gradient_accumulation_steps=args.gradient_accumulation_steps,
+#         fp16=args.fp16,
+#         bf16=args.bf16,
+#         dataloader_num_workers=args.num_workers,
+#         remove_unused_columns=False,
+#         seed=args.seed,
+#     )
+    
+#     print("Initializing trainer...")
+#     # Initialize the contrastive trainer
+#     trainer = ContrastiveTrainer(
+#         model=model,
+#         args=training_args,
+#         train_dataset=train_dataset,
+#         eval_dataset=eval_dataset,
+#         data_collator=protein_text_collate_fn,
+#     )
+    
+#     print("Starting contrastive training...")
+#     # Train the model
+#     trainer.train()
+    
+#     print("Saving final model...")
+#     # Save the final model
+#     trainer.save_model()
+    
+#     # Save only the projection layer weights for later use
+#     projection_path = os.path.join(args.output_dir, "protein_projection.pt")
+#     torch.save(model.protein_projection.state_dict(), projection_path)
+#     print(f"Saved protein projection weights to: {projection_path}")
+    
+#     # Final evaluation
+#     if eval_dataset:
+#         print("Running final evaluation...")
+#         eval_results = trainer.evaluate()
+#         print(f"Final evaluation results: {eval_results}")
+        
+#         if args.use_wandb:
+#             wandb.log({"final_eval": eval_results})
+    
+#     print("Contrastive training completed!")
+    
+#     if args.use_wandb:
+#         wandb.finish()
+    
+#     return trainer
+
+
+# if __name__ == "__main__":
+#     parser = ArgumentParser(description="Contrastive pre-training for Protein-LLM")
+    
+#     # Model configuration
+#     parser.add_argument("--text_model_name", type=str, default="Qwen/Qwen3-1.7B",
+#                         help="Name or path to the text model")
+#     parser.add_argument("--protein_model_name", type=str, default="facebook/esm2_t6_8M_UR50D",
+#                         help="Name or path to the protein model")
+#     parser.add_argument("--qformer_model_name", type=str, 
+#                         default="microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext",
+#                         help="Name or path to the QFormer model")
+#     parser.add_argument("--cache_dir", type=str, default="/model-weights",
+#                         help="Directory to cache downloaded models")
+#     parser.add_argument("--num_query_tokens", type=int, default=32,
+#                         help="Number of query tokens in QFormer")
+    
+#     # Dataset configuration
+#     parser.add_argument("--dataset_name", type=str, default="wanglab/protein_descriptions",
+#                         help="Name of the dataset for contrastive learning")
+#     parser.add_argument("--eval_dataset", action="store_true",
+#                         help="Whether to use evaluation dataset")
+    
+#     # Training configuration
+#     parser.add_argument("--output_dir", type=str, default="./contrastive_outputs",
+#                         help="Output directory for model and logs")
+#     parser.add_argument("--num_epochs", type=int, default=10,
+#                         help="Number of training epochs")
+#     parser.add_argument("--batch_size", type=int, default=32,
+#                         help="Batch size per device")
+#     parser.add_argument("--learning_rate", type=float, default=1e-4,
+#                         help="Learning rate")
+#     parser.add_argument("--weight_decay", type=float, default=0.01,
+#                         help="Weight decay")
+#     parser.add_argument("--warmup_steps", type=int, default=1000,
+#                         help="Number of warmup steps")
+#     parser.add_argument("--gradient_accumulation_steps", type=int, default=1,
+#                         help="Gradient accumulation steps")
+    
+#     # Contrastive learning specific
+#     parser.add_argument("--temperature", type=float, default=0.07,
+#                         help="Temperature for contrastive loss")
+#     parser.add_argument("--freeze_protein_model", action="store_true", default=True,
+#                         help="Freeze protein model during training")
+#     parser.add_argument("--freeze_text_model", action="store_true", default=True,
+#                         help="Freeze text model during training")
+#     parser.add_argument("--protein_weight", type=float, default=1.0,
+#                         help="Weight for protein features in contrastive loss")
+#     parser.add_argument("--text_weight", type=float, default=1.0,
+#                         help="Weight for text features in contrastive loss")
+    
+#     # Data configuration
+#     parser.add_argument("--max_length_protein", type=int, default=1024,
+#                         help="Maximum length for protein sequences")
+#     parser.add_argument("--max_length_text", type=int, default=512,
+#                         help="Maximum length for text sequences")
+#     parser.add_argument("--num_workers", type=int, default=4,
+#                         help="Number of data loading workers")
+    
+#     # Logging and evaluation
+#     parser.add_argument("--logging_steps", type=int, default=100,
+#                         help="Number of steps between logging")
+#     parser.add_argument("--eval_steps", type=int, default=500,
+#                         help="Number of steps between evaluations")
+#     parser.add_argument("--save_steps", type=int, default=1000,
+#                         help="Number of steps between saving checkpoints")
+#     parser.add_argument("--save_total_limit", type=int, default=3,
+#                         help="Maximum number of checkpoints to keep")
+    
+#     # Hardware configuration
+#     parser.add_argument("--fp16", action="store_true",
+#                         help="Use FP16 precision")
+#     parser.add_argument("--bf16", action="store_true",
+#                         help="Use BF16 precision")
+#     parser.add_argument("--seed", type=int, default=42,
+#                         help="Random seed")
+    
+#     # Wandb logging
+#     parser.add_argument("--use_wandb", action="store_true",
+#                         help="Use Weights & Biases for logging")
+#     parser.add_argument("--wandb_project", type=str, default="protein-llm-contrastive",
+#                         help="Wandb project name")
+#     parser.add_argument("--wandb_entity", type=str, default=None,
+#                         help="Wandb entity name")
+    
+#     args = parser.parse_args()
+    
+#     # Create output directory
+#     os.makedirs(args.output_dir, exist_ok=True)
+    
+#     # Run contrastive training
+#     trainer = main(args)
+
+
+import os
+import time
+from argparse import ArgumentParser
+from functools import partial
+
+import torch
+import wandb
+from datasets import load_dataset
+from torch.utils.data import DataLoader
+
+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.models.protein_llm import ProteinLLMModel
+from bioreason.trainer.contrast_trainer_new import (
+    ContrastiveTrainer,
+    ContrastiveTrainingArguments,
+    protein_text_collate_fn,
+)
+from bioreason.dataset.protein import format_protein_contrastive
+
+
+def main(args):
+    """
+    Main function for enhanced contrastive pre-training of Protein-LLM.
+    
+    This script trains the QFormer projection layer to align protein and text representations
+    using enhanced contrastive learning with optional protein-text matching before fine-tuning.
+    """
+    
+    # Set random seed
+    pl.seed_everything(args.seed)
+    torch.cuda.empty_cache()
+    torch.set_float32_matmul_precision("medium")
+    
+    # Initialize wandb
+    if args.use_wandb:
+        wandb.init(
+            project=args.wandb_project,
+            entity=args.wandb_entity,
+            name=f"enhanced-contrastive-{args.text_model_name.split('/')[-1]}-{time.strftime('%Y%m%d-%H%M%S')}",
+            config=vars(args)
+        )
+    
+    print("Loading model...")
+    # Load the Protein-LLM model
+    model = ProteinLLMModel(
+        text_model_name=args.text_model_name,
+        protein_model_name=args.protein_model_name,
+        qformer_model_name=args.qformer_model_name,
+        cache_dir=args.cache_dir,
+        max_length_protein=args.max_length_protein,
+        max_length_text=args.max_length_text,
+        text_model_finetune=False,  # Don't fine-tune during contrastive learning
+        protein_model_finetune=False,  # Don't fine-tune during contrastive learning
+        num_query_tokens=args.num_query_tokens,
+    )
+    
+    print("Loading datasets...")
+    # Load datasets for contrastive learning
+    train_dataset = load_dataset("json", data_files=args.train_dataset, split="train")
+    eval_dataset = load_dataset("json", data_files=args.valid_dataset, split="train") if args.eval_dataset else None
+    
+    # Format datasets for contrastive learning
+    train_dataset = train_dataset.map(format_protein_contrastive)
+    if eval_dataset:
+        eval_dataset = eval_dataset.map(format_protein_contrastive)
+    
+    # Filter out examples without protein sequences or descriptions
+    train_dataset = train_dataset.filter(
+        lambda x: x["protein"] and x["text"] 
+        and len(x["protein"].strip()) > 0 
+        and len(x["text"].strip()) > 0
+    )
+    
+    if eval_dataset:
+        eval_dataset = eval_dataset.filter(
+            lambda x: x["protein"] and x["text"] 
+            and len(x["protein"].strip()) > 0 
+            and len(x["text"].strip()) > 0
+        )
+    
+    print(f"Training dataset size: {len(train_dataset)}")
+    if eval_dataset:
+        print(f"Eval dataset size: {len(eval_dataset)}")
+    
+    # Setup enhanced training arguments for contrastive learning
+    
+    training_args = ContrastiveTrainingArguments(
+        output_dir=args.output_dir,
+        num_train_epochs=args.num_epochs,
+        per_device_train_batch_size=args.batch_size,
+        per_device_eval_batch_size=args.batch_size,
+        learning_rate=args.learning_rate,
+        weight_decay=args.weight_decay,
+        temperature=args.temperature,
+        freeze_protein_model=args.freeze_protein_model,
+        freeze_text_model=args.freeze_text_model,
+        protein_weight=args.protein_weight,
+        text_weight=args.text_weight,
+        enable_ptm=args.enable_ptm,
+        ptm_weight=args.ptm_weight,
+        max_length_protein=args.max_length_protein,
+        max_length_text=args.max_length_text,
+        logging_steps=args.logging_steps,
+        eval_strategy="steps" if eval_dataset else "no",
+        eval_steps=args.eval_steps if eval_dataset else None,
+        save_steps=args.save_steps,
+        save_total_limit=args.save_total_limit,
+        load_best_model_at_end=True if eval_dataset else False,
+        metric_for_best_model="eval_avg_recall_at_1" if eval_dataset else None,
+        greater_is_better=True,
+        report_to=["wandb"] if args.use_wandb else [],
+        warmup_steps=args.warmup_steps,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        fp16=args.fp16,
+        bf16=args.bf16,
+        dataloader_num_workers=args.num_workers,
+        remove_unused_columns=False,
+        seed=args.seed,
+        # Distributed training settings
+        ddp_find_unused_parameters=False,
+        dataloader_pin_memory=True,
+    )
+    
+    print("Initializing enhanced trainer...")
+    # Initialize the enhanced contrastive trainer
+    trainer = ContrastiveTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        data_collator=protein_text_collate_fn,
+    )
+    
+    print("Starting enhanced contrastive training...")
+    print(f"- Contrastive learning enabled")
+    print(f"- Protein-text matching: {'enabled' if args.enable_ptm else 'disabled'}")
+    print(f"- Temperature: {args.temperature}")
+    print(f"- PTM weight: {args.ptm_weight}")
+    
+    # Train the model
+    trainer.train()
+    
+    print("Saving final model...")
+    # Save the final model
+    trainer.save_model()
+    
+    # Save projection layer weights and PTM head if enabled
+    projection_path = os.path.join(args.output_dir, "protein_projection.pt")
+    torch.save(model.protein_projection.state_dict(), projection_path)
+    print(f"Saved protein projection weights to: {projection_path}")
+    
+    if args.enable_ptm and hasattr(trainer.contrastive_loss, 'ptm_head'):
+        ptm_head_path = os.path.join(args.output_dir, "ptm_head.pt")
+        torch.save(trainer.contrastive_loss.ptm_head.state_dict(), ptm_head_path)
+        print(f"Saved PTM head weights to: {ptm_head_path}")
+    
+    # Final evaluation
+    if eval_dataset:
+        print("Running final evaluation...")
+        eval_results = trainer.evaluate()
+        print(f"Final evaluation results: {eval_results}")
+        
+        if args.use_wandb:
+            wandb.log({"final_eval": eval_results})
+    
+    print("Enhanced contrastive training completed!")
+    
+    if args.use_wandb:
+        wandb.finish()
+    
+    return trainer
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser(description="Enhanced contrastive pre-training for Protein-LLM")
+    
+    # Model configuration
+    parser.add_argument("--text_model_name", type=str, default="Qwen/Qwen3-1.7B",
+                        help="Name or path to the text model")
+    parser.add_argument("--protein_model_name", type=str, default="facebook/esm2_t6_8M_UR50D",
+                        help="Name or path to the protein model")
+    parser.add_argument("--qformer_model_name", type=str, 
+                        default="microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext",
+                        help="Name or path to the QFormer model")
+    parser.add_argument("--cache_dir", type=str, default="/model-weights",
+                        help="Directory to cache downloaded models")
+    parser.add_argument("--num_query_tokens", type=int, default=32,
+                        help="Number of query tokens in QFormer")
+    
+    # Dataset configuration
+    parser.add_argument("--train_dataset", type=str, default="wanglab/protein_descriptions",
+                        help="Name of the dataset for contrastive learning")
+    parser.add_argument("--valid_dataset", type=str, default="wanglab/protein_descriptions",
+                        help="Name of the dataset for contrastive learning")
+    parser.add_argument("--eval_dataset", action="store_true",
+                        help="Whether to use evaluation dataset")
+    
+    # Training configuration
+    parser.add_argument("--output_dir", type=str, default="./enhanced_contrastive_outputs",
+                        help="Output directory for model and logs")
+    parser.add_argument("--num_epochs", type=int, default=10,
+                        help="Number of training epochs")
+    parser.add_argument("--batch_size", type=int, default=32,
+                        help="Batch size per device")
+    parser.add_argument("--learning_rate", type=float, default=1e-4,
+                        help="Learning rate")
+    parser.add_argument("--weight_decay", type=float, default=0.01,
+                        help="Weight decay")
+    parser.add_argument("--warmup_steps", type=int, default=1000,
+                        help="Number of warmup steps")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=1,
+                        help="Gradient accumulation steps")
+    
+    # Enhanced contrastive learning specific
+    parser.add_argument("--temperature", type=float, default=0.07,
+                        help="Temperature for contrastive loss")
+    parser.add_argument("--freeze_protein_model", action="store_true", default=True,
+                        help="Freeze protein model during training")
+    parser.add_argument("--freeze_text_model", action="store_true", default=True,
+                        help="Freeze text model during training")
+    parser.add_argument("--protein_weight", type=float, default=1.0,
+                        help="Weight for protein features in contrastive loss")
+    parser.add_argument("--text_weight", type=float, default=1.0,
+                        help="Weight for text features in contrastive loss")
+    
+    # Protein-Text Matching (PTM) configuration
+    parser.add_argument("--enable_ptm", action="store_true", default=True,
+                        help="Enable protein-text matching task")
+    parser.add_argument("--ptm_weight", type=float, default=1.0,
+                        help="Weight for protein-text matching loss")
+    
+    # Data configuration
+    parser.add_argument("--max_length_protein", type=int, default=1024,
+                        help="Maximum length for protein sequences")
+    parser.add_argument("--max_length_text", type=int, default=512,
+                        help="Maximum length for text sequences")
+    parser.add_argument("--num_workers", type=int, default=4,
+                        help="Number of data loading workers")
+    
+    # Logging and evaluation
+    parser.add_argument("--logging_steps", type=int, default=100,
+                        help="Number of steps between logging")
+    parser.add_argument("--eval_steps", type=int, default=500,
+                        help="Number of steps between evaluations")
+    parser.add_argument("--save_steps", type=int, default=1000,
+                        help="Number of steps between saving checkpoints")
+    parser.add_argument("--save_total_limit", type=int, default=3,
+                        help="Maximum number of checkpoints to keep")
+    
+    # Hardware configuration
+    parser.add_argument("--fp16", action="store_true",
+                        help="Use FP16 precision")
+    parser.add_argument("--bf16", action="store_true",
+                        help="Use BF16 precision")
+    parser.add_argument("--seed", type=int, default=42,
+                        help="Random seed")
+    
+    # Wandb logging
+    parser.add_argument("--use_wandb", action="store_true",
+                        help="Use Weights & Biases for logging")
+    parser.add_argument("--wandb_project", type=str, default="protein-llm-enhanced-contrastive",
+                        help="Wandb project name")
+    parser.add_argument("--wandb_entity", type=str, default=None,
+                        help="Wandb entity name")
+    
+    args = parser.parse_args()
+    
+    # Validate arguments
+    if args.enable_ptm and not hasattr(args, 'ptm_weight'):
+        args.ptm_weight = 1.0
+    
+    # Create output directory
+    os.makedirs(args.output_dir, exist_ok=True)
+    
+    # Print configuration
+    print("=" * 50)
+    print("Enhanced Contrastive Training Configuration:")
+    print("=" * 50)
+    print(f"Text model: {args.text_model_name}")
+    print(f"Protein model: {args.protein_model_name}")
+    print(f"QFormer model: {args.qformer_model_name}")
+    print(f"Dataset: {args.train_dataset}")
+    print(f"Output directory: {args.output_dir}")
+    print(f"Batch size: {args.batch_size}")
+    print(f"Learning rate: {args.learning_rate}")
+    print(f"Temperature: {args.temperature}")
+    print(f"Enable PTM: {args.enable_ptm}")
+    print(f"PTM weight: {args.ptm_weight}")
+    print(f"Number of epochs: {args.num_epochs}")
+    print(f"Query tokens: {args.num_query_tokens}")
+    print("=" * 50)
+    
+    # Run enhanced contrastive training
+    trainer = main(args)

BioReason_new/train_protein_qwen.py

@@ -0,0 +1,839 @@
+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.protein import get_format_protein_function, protein_llm_collate_fn
+from bioreason.dataset.utils import truncate_protein
+from bioreason.models.dl.processing_dl import ProteinLLMProcessor
+from bioreason.models.protein_llm import ProteinLLMModel
+
+# Set start method to 'spawn' for CUDA compatibility with multiprocessing
+torch.multiprocessing.set_sharing_strategy("file_system")
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+class ProteinLLMFineTuner(pl.LightningModule):
+    """
+    PyTorch Lightning module for fine-tuning Protein-LLM models.
+    """
+
+    def __init__(self, hparams):
+        """
+        Initialize the ProteinLLMFineTuner.
+
+        Args:
+            hparams: Hyperparameters for the model and training
+        """
+        super().__init__()
+        self.save_hyperparameters(hparams)
+
+        self.text_model_name = self.hparams.text_model_name
+        self.protein_model_name = self.hparams.protein_model_name
+        self.qformer_model_name = self.hparams.qformer_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.protein_model_finetune = self.hparams.protein_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_protein = self.hparams.max_length_protein
+        self.max_length_text = self.hparams.max_length_text
+        self.num_query_tokens = self.hparams.num_query_tokens
+        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 = ProteinLLMModel(
+            text_model_name=self.text_model_name,
+            protein_model_name=self.protein_model_name,
+            qformer_model_name=self.qformer_model_name,
+            cache_dir=self.cache_dir,
+            max_length_protein=self.max_length_protein,
+            max_length_text=self.max_length_text,
+            text_model_finetune=self.text_model_finetune,
+            protein_model_finetune=self.protein_model_finetune,
+            num_query_tokens=self.num_query_tokens,
+        )
+
+        self.text_model = self.model.text_model
+        self.protein_model = self.model.protein_model
+        self.protein_projection = self.model.protein_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 ProteinLLMModel.
+        """
+
+        # Freeze protein encoder parameters
+        if self.protein_model_finetune:
+            pass
+        else:
+            for param in self.protein_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
+            lora_config = None
+
+        # Make projection layer trainable
+        for param in self.protein_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
+        protein_tokenized = batch.get("protein_tokenized")
+        if protein_tokenized is not None:
+            protein_tokenized = protein_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,
+            protein_tokenized=protein_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
+        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 protein data for this example
+                    example_protein_data = None
+                    example_batch_map = None
+
+                    if protein_tokenized is not None and batch_idx_map is not None:
+                        # Find protein 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 protein data
+                            example_protein_data = BatchEncoding(
+                                {
+                                    "input_ids": protein_tokenized.input_ids[example_indices].to(self.device),
+                                    "attention_mask": protein_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,
+                            protein_tokenized=example_protein_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_protein_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 == "protein_function":
+            # Use Hugging Face dataset if provided
+            dataset = load_dataset(self.hparams.protein_function_data_dir_huggingface)
+            dataset = dataset.map(get_format_protein_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_protein_per_side:
+                train_dataset = train_dataset.map(
+                    truncate_protein, fn_kwargs={"truncate_protein_per_side": self.hparams.truncate_protein_per_side}
+                )
+            
+            processor = ProteinLLMProcessor(
+                tokenizer=self.model.text_tokenizer,
+                protein_tokenizer=self.model.protein_tokenizer,
+            )
+
+            # Create partial function with all required arguments except the batch
+            collate_fn = partial(
+                protein_llm_collate_fn,
+                processor=processor,
+                max_length_text=self.max_length_text,
+                max_length_protein=self.max_length_protein,
+                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 == "protein_function":
+            # Use Hugging Face dataset
+            dataset = load_dataset(self.hparams.protein_function_data_dir_huggingface)
+            dataset = dataset.map(get_format_protein_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_protein_per_side:
+                val_dataset = val_dataset.map(
+                    truncate_protein, fn_kwargs={"truncate_protein_per_side": self.hparams.truncate_protein_per_side}
+                )
+
+            processor = ProteinLLMProcessor(
+                tokenizer=self.model.text_tokenizer,
+                protein_tokenizer=self.model.protein_tokenizer,
+            )
+
+            # Create partial function with all required arguments except the batch
+            collate_fn = partial(
+                protein_llm_collate_fn,
+                processor=processor,
+                max_length_text=self.max_length_text,
+                max_length_protein=self.max_length_protein,
+                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()
+    
+    # For protein function datasets, 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 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] if len(self.labels) > 0 else "negative"
+        pos_label = self.labels[1] if len(self.labels) > 1 else "positive"
+        
+        # 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
+        correct_predictions = 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"]
+            protein_tokenized = batch.get("protein_tokenized")
+            if protein_tokenized is not None:
+                protein_tokenized = protein_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)
+            
+            # Process examples in the batch
+            examples_in_batch = 0
+            for example_idx in range(input_ids.size(0)):
+                # 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
+                
+                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 protein data for this example
+                    example_protein_data = None
+                    example_batch_map = None
+                    
+                    if protein_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_protein_data = BatchEncoding({
+                                "input_ids": protein_tokenized.input_ids[example_indices].to(self.device),
+                                "attention_mask": protein_tokenized.attention_mask[example_indices].to(self.device),
+                            })
+                            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,
+                            protein_tokenized=example_protein_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]
+                    
+                    # Check if the generated text contains the ground truth
+                    generation_contains_ground_truth = ground_truth.lower() in generation.lower()
+                    
+                    # Update metrics
+                    total_examples += 1
+                    examples_in_batch += 1
+                    
+                    if generation_contains_ground_truth:
+                        correct_predictions += 1
+                    
+                    # 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,
+                    })
+                    
+                    # Clean up memory
+                    torch.cuda.empty_cache()
+                    gc.collect()
+            
+            # Log batch completion to wandb
+            processed_batches += 1
+            current_accuracy = correct_predictions / max(total_examples, 1)
+            
+            wandb_logger.log({
+                "batches_processed": processed_batches,
+                "examples_processed": total_examples,
+                "examples_in_last_batch": examples_in_batch,
+                "current_accuracy": current_accuracy,
+                "progress_percentage": (batch_idx + 1) / total_batches * 100
+            })
+        
+        # Calculate final metrics
+        accuracy = correct_predictions / max(total_examples, 1)
+        
+        # Log final metrics to wandb
+        wandb_logger.log({
+            "test_accuracy": accuracy,
+            "correct_predictions": correct_predictions,
+            "total_examples_processed": total_examples,
+            "test_status": "completed"
+        })
+        
+        # Create a table with all the generations
+        if generations:
+            columns = [
+                "batch_idx",
+                "example_idx", 
+                "user_input",
+                "generation",
+                "ground_truth",
+                "contains_ground_truth"
+            ]
+            data = []
+            for g in generations:
+                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"Correct: {correct_predictions}\n"
+        )
+        print(summary)
+        wandb_logger.log({"test_summary": summary})
+        
+        # Force garbage collection
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        return {
+            "test_accuracy": accuracy,
+        }
+
+
+def main(args: ArgumentParser):
+    """
+    Main function to run the Protein-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 = ProteinLLMFineTuner(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,
+    )
+
+    # 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", "protein-llm"], default="protein-llm")
+    parser.add_argument("--text_model_name", type=str, default="Qwen/Qwen3-1.7B")
+    parser.add_argument("--protein_model_name", type=str, default="facebook/esm2_t6_8M_UR50D")
+    parser.add_argument("--qformer_model_name", type=str, default="microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext")
+    parser.add_argument("--text_model_finetune", type=bool, default=True)
+    parser.add_argument("--protein_model_finetune", type=bool, default=False)
+    parser.add_argument("--num_query_tokens", type=int, default=32)
+    
+    # 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_protein", type=int, default=1024)
+    parser.add_argument("--max_length_text", type=int, default=1024)
+    parser.add_argument("--truncate_protein_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=["protein_function"], default="protein_function")
+    parser.add_argument("--use_protein_llm_collate_fn", type=bool, default=True)
+    parser.add_argument("--protein_function_data_dir_huggingface", type=str, default="wanglab/protein_function")
+    parser.add_argument("--merge_val_test_set", type=bool, default=False)
+    
+    # Logging and monitoring
+    parser.add_argument("--wandb_project", type=str, default="esm2-qwen3-1.7b-finetune")
+    parser.add_argument("--wandb_entity", type=str)
+    
+    args = parser.parse_args()
+
+    main(args)

BioReason_new/wandb/debug-internal.log

@@ -0,0 +1,28 @@
+{"time":"2025-08-13T15:33:18.937087932+08:00","level":"INFO","msg":"stream: starting","core version":"0.21.1"}
+{"time":"2025-08-13T15:33:22.125246419+08:00","level":"INFO","msg":"stream: created new stream","id":"ig4rhoqf"}
+{"time":"2025-08-13T15:33:22.126020019+08:00","level":"INFO","msg":"stream: started","id":"ig4rhoqf"}
+{"time":"2025-08-13T15:33:22.12605541+08:00","level":"INFO","msg":"writer: started","stream_id":"ig4rhoqf"}
+{"time":"2025-08-13T15:33:22.126066944+08:00","level":"INFO","msg":"handler: started","stream_id":"ig4rhoqf"}
+{"time":"2025-08-13T15:33:22.126093203+08:00","level":"INFO","msg":"sender: started","stream_id":"ig4rhoqf"}
+{"time":"2025-08-13T15:33:29.266636932+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=1c9fbc6cbd51ae71160b439aa8a692f95ee16291fd45e2c33912f90a5ae3a56a09f33e745a416b04da501c1cf40827caa3399f6cde0752ce3f571c12353aa477d27f258d11b74762a029b2fada9797f27a85c7fe03e3bb1ec12fb98789a6b389cf000b2742ca658544ea41f22c9631272e67f2ea275c4a5b7b9c516fed6ed22ccce6b26b684c341e337963a684be9b248bd64b665fb73437d28ce7365767fe6721770c9f8f75e150767d9e4154a983b7f6aafcc31bbc050a26df2e9a3b99e065c111d7a92525568b04cf8d8a7490e6b8f77aed12c26aa378d0708385b0f40b22af3988b5c6cfb25918b1c9471be072bdfdc7a986ca0774f732d27361d23f3448&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:52222->142.250.73.123:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=1c9fbc6cbd51ae71160b439aa8a692f95ee16291fd45e2c33912f90a5ae3a56a09f33e745a416b04da501c1cf40827caa3399f6cde0752ce3f571c12353aa477d27f258d11b74762a029b2fada9797f27a85c7fe03e3bb1ec12fb98789a6b389cf000b2742ca658544ea41f22c9631272e67f2ea275c4a5b7b9c516fed6ed22ccce6b26b684c341e337963a684be9b248bd64b665fb73437d28ce7365767fe6721770c9f8f75e150767d9e4154a983b7f6aafcc31bbc050a26df2e9a3b99e065c111d7a92525568b04cf8d8a7490e6b8f77aed12c26aa378d0708385b0f40b22af3988b5c6cfb25918b1c9471be072bdfdc7a986ca0774f732d27361d23f3448&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:33:31.205628552+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/requirements.txt?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=27041959870a9e2ba8f974ca28bb5723ad374fc1917b982522ea3f555344b25fed8f73d5b443663eb1d7d10499d748efa90c8f6d91802daabf9e8c31258a7371e93aca1694d3e8bcb35a72d7c9ca243dc164f5dfae6ec3009c25cc78ea8cf37629f017b7538998f44f7c65ccfb675e343601421475a4490c754c5ee0370d3d5dfa928faddfe9a90621302ce69efd3d26c51f49c23f92148b018281ccd02f22c42e73e318594ea9c2ff9b25ad13163b60c37f1ededb4dc3a50712bcabdffed71883e2e5b5f04c40f13612d9ff1b6762b9f79d6e19873e959fbf9495e4827e901e8dbd25d7b6f291841915bc428212ad0a142cc3a9b7d04dbba9b49873573481fa&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:57666->142.250.73.91:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/requirements.txt?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=27041959870a9e2ba8f974ca28bb5723ad374fc1917b982522ea3f555344b25fed8f73d5b443663eb1d7d10499d748efa90c8f6d91802daabf9e8c31258a7371e93aca1694d3e8bcb35a72d7c9ca243dc164f5dfae6ec3009c25cc78ea8cf37629f017b7538998f44f7c65ccfb675e343601421475a4490c754c5ee0370d3d5dfa928faddfe9a90621302ce69efd3d26c51f49c23f92148b018281ccd02f22c42e73e318594ea9c2ff9b25ad13163b60c37f1ededb4dc3a50712bcabdffed71883e2e5b5f04c40f13612d9ff1b6762b9f79d6e19873e959fbf9495e4827e901e8dbd25d7b6f291841915bc428212ad0a142cc3a9b7d04dbba9b49873573481fa&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:33:31.959506763+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=1c9fbc6cbd51ae71160b439aa8a692f95ee16291fd45e2c33912f90a5ae3a56a09f33e745a416b04da501c1cf40827caa3399f6cde0752ce3f571c12353aa477d27f258d11b74762a029b2fada9797f27a85c7fe03e3bb1ec12fb98789a6b389cf000b2742ca658544ea41f22c9631272e67f2ea275c4a5b7b9c516fed6ed22ccce6b26b684c341e337963a684be9b248bd64b665fb73437d28ce7365767fe6721770c9f8f75e150767d9e4154a983b7f6aafcc31bbc050a26df2e9a3b99e065c111d7a92525568b04cf8d8a7490e6b8f77aed12c26aa378d0708385b0f40b22af3988b5c6cfb25918b1c9471be072bdfdc7a986ca0774f732d27361d23f3448&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:52232->142.250.73.123:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=1c9fbc6cbd51ae71160b439aa8a692f95ee16291fd45e2c33912f90a5ae3a56a09f33e745a416b04da501c1cf40827caa3399f6cde0752ce3f571c12353aa477d27f258d11b74762a029b2fada9797f27a85c7fe03e3bb1ec12fb98789a6b389cf000b2742ca658544ea41f22c9631272e67f2ea275c4a5b7b9c516fed6ed22ccce6b26b684c341e337963a684be9b248bd64b665fb73437d28ce7365767fe6721770c9f8f75e150767d9e4154a983b7f6aafcc31bbc050a26df2e9a3b99e065c111d7a92525568b04cf8d8a7490e6b8f77aed12c26aa378d0708385b0f40b22af3988b5c6cfb25918b1c9471be072bdfdc7a986ca0774f732d27361d23f3448&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:33:33.871387522+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/requirements.txt?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=27041959870a9e2ba8f974ca28bb5723ad374fc1917b982522ea3f555344b25fed8f73d5b443663eb1d7d10499d748efa90c8f6d91802daabf9e8c31258a7371e93aca1694d3e8bcb35a72d7c9ca243dc164f5dfae6ec3009c25cc78ea8cf37629f017b7538998f44f7c65ccfb675e343601421475a4490c754c5ee0370d3d5dfa928faddfe9a90621302ce69efd3d26c51f49c23f92148b018281ccd02f22c42e73e318594ea9c2ff9b25ad13163b60c37f1ededb4dc3a50712bcabdffed71883e2e5b5f04c40f13612d9ff1b6762b9f79d6e19873e959fbf9495e4827e901e8dbd25d7b6f291841915bc428212ad0a142cc3a9b7d04dbba9b49873573481fa&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:52238->142.250.73.123:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/requirements.txt?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=27041959870a9e2ba8f974ca28bb5723ad374fc1917b982522ea3f555344b25fed8f73d5b443663eb1d7d10499d748efa90c8f6d91802daabf9e8c31258a7371e93aca1694d3e8bcb35a72d7c9ca243dc164f5dfae6ec3009c25cc78ea8cf37629f017b7538998f44f7c65ccfb675e343601421475a4490c754c5ee0370d3d5dfa928faddfe9a90621302ce69efd3d26c51f49c23f92148b018281ccd02f22c42e73e318594ea9c2ff9b25ad13163b60c37f1ededb4dc3a50712bcabdffed71883e2e5b5f04c40f13612d9ff1b6762b9f79d6e19873e959fbf9495e4827e901e8dbd25d7b6f291841915bc428212ad0a142cc3a9b7d04dbba9b49873573481fa&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:33:36.261477525+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=1c9fbc6cbd51ae71160b439aa8a692f95ee16291fd45e2c33912f90a5ae3a56a09f33e745a416b04da501c1cf40827caa3399f6cde0752ce3f571c12353aa477d27f258d11b74762a029b2fada9797f27a85c7fe03e3bb1ec12fb98789a6b389cf000b2742ca658544ea41f22c9631272e67f2ea275c4a5b7b9c516fed6ed22ccce6b26b684c341e337963a684be9b248bd64b665fb73437d28ce7365767fe6721770c9f8f75e150767d9e4154a983b7f6aafcc31bbc050a26df2e9a3b99e065c111d7a92525568b04cf8d8a7490e6b8f77aed12c26aa378d0708385b0f40b22af3988b5c6cfb25918b1c9471be072bdfdc7a986ca0774f732d27361d23f3448&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:57678->142.250.73.91:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=1c9fbc6cbd51ae71160b439aa8a692f95ee16291fd45e2c33912f90a5ae3a56a09f33e745a416b04da501c1cf40827caa3399f6cde0752ce3f571c12353aa477d27f258d11b74762a029b2fada9797f27a85c7fe03e3bb1ec12fb98789a6b389cf000b2742ca658544ea41f22c9631272e67f2ea275c4a5b7b9c516fed6ed22ccce6b26b684c341e337963a684be9b248bd64b665fb73437d28ce7365767fe6721770c9f8f75e150767d9e4154a983b7f6aafcc31bbc050a26df2e9a3b99e065c111d7a92525568b04cf8d8a7490e6b8f77aed12c26aa378d0708385b0f40b22af3988b5c6cfb25918b1c9471be072bdfdc7a986ca0774f732d27361d23f3448&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:33:41.039143093+08:00","level":"INFO","msg":"stream: closing","id":"ig4rhoqf"}
+{"time":"2025-08-13T15:38:00.686011134+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/requirements.txt?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=27041959870a9e2ba8f974ca28bb5723ad374fc1917b982522ea3f555344b25fed8f73d5b443663eb1d7d10499d748efa90c8f6d91802daabf9e8c31258a7371e93aca1694d3e8bcb35a72d7c9ca243dc164f5dfae6ec3009c25cc78ea8cf37629f017b7538998f44f7c65ccfb675e343601421475a4490c754c5ee0370d3d5dfa928faddfe9a90621302ce69efd3d26c51f49c23f92148b018281ccd02f22c42e73e318594ea9c2ff9b25ad13163b60c37f1ededb4dc3a50712bcabdffed71883e2e5b5f04c40f13612d9ff1b6762b9f79d6e19873e959fbf9495e4827e901e8dbd25d7b6f291841915bc428212ad0a142cc3a9b7d04dbba9b49873573481fa&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:44768->142.250.217.91:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/requirements.txt?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=27041959870a9e2ba8f974ca28bb5723ad374fc1917b982522ea3f555344b25fed8f73d5b443663eb1d7d10499d748efa90c8f6d91802daabf9e8c31258a7371e93aca1694d3e8bcb35a72d7c9ca243dc164f5dfae6ec3009c25cc78ea8cf37629f017b7538998f44f7c65ccfb675e343601421475a4490c754c5ee0370d3d5dfa928faddfe9a90621302ce69efd3d26c51f49c23f92148b018281ccd02f22c42e73e318594ea9c2ff9b25ad13163b60c37f1ededb4dc3a50712bcabdffed71883e2e5b5f04c40f13612d9ff1b6762b9f79d6e19873e959fbf9495e4827e901e8dbd25d7b6f291841915bc428212ad0a142cc3a9b7d04dbba9b49873573481fa&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:38:04.732136338+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/output.log?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=1d3cd94b4887f4c0adaa644487d39281fc878818fa20fa0c39b307b4d29a0774e1077acff64b5047465d575c9417f96645240e67e6dab6d1cc25f617e622bcf0f8715629cc8509684074bcd0688360af452c0f89dc80a4e9025500a6e0d16485ec7617abfd0ff10c8ae62cf6012c0f43aad8d30bb7c9de79f87c56f73ccabc572e83285debd09c51e5fcc95e1a7fe1cc011606044e056dc5f2d6756e365041194a3e52cc500d0e71499406481acf9cc760c1da018d07a35e07830dd04208781140f1fc38c86f5aa24bf65e1a58d35406ae84e7026fe6cf59353cf16ca53b51a813beb22417486e4fabc77de1aaef6ee20b755ce21b765571f6bfbf8473942f9d&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:44778->142.250.217.91:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/output.log?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=1d3cd94b4887f4c0adaa644487d39281fc878818fa20fa0c39b307b4d29a0774e1077acff64b5047465d575c9417f96645240e67e6dab6d1cc25f617e622bcf0f8715629cc8509684074bcd0688360af452c0f89dc80a4e9025500a6e0d16485ec7617abfd0ff10c8ae62cf6012c0f43aad8d30bb7c9de79f87c56f73ccabc572e83285debd09c51e5fcc95e1a7fe1cc011606044e056dc5f2d6756e365041194a3e52cc500d0e71499406481acf9cc760c1da018d07a35e07830dd04208781140f1fc38c86f5aa24bf65e1a58d35406ae84e7026fe6cf59353cf16ca53b51a813beb22417486e4fabc77de1aaef6ee20b755ce21b765571f6bfbf8473942f9d&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:38:04.732752873+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-summary.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=43378ba7835541b25e7884ddd1295c99cdb59af218b575e11ebf49219cd7b70414883b2eaabf6ce5bf29b6dcf3aa4854cc71ac3f00aceb098b2bd9f008cbf3c3e776aeb0eec09eafbd1633dbdf5cd30ca267e7fc78a69846132727fb83834b13e3aaf2e17cf6af4974dc4f96f982c762fbd056b369608aece84f3c97eed6c02b7648d28a1ab3f3566026c830a6eb4d57a55676a251768961021c9050a7b6974e74c7efd13ca9324ac1a6768bca2daa0daf12b3c8ebc3ad02e5d5c935024048da77f0b6047842ebceaf42ec881dd23660f5f4caad61c04ea8736f4b8abd5c81197ebd10853a45e4944b85d318500c385596344922b43e114698bc05474aad582c&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:43548->142.250.73.123:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-summary.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=43378ba7835541b25e7884ddd1295c99cdb59af218b575e11ebf49219cd7b70414883b2eaabf6ce5bf29b6dcf3aa4854cc71ac3f00aceb098b2bd9f008cbf3c3e776aeb0eec09eafbd1633dbdf5cd30ca267e7fc78a69846132727fb83834b13e3aaf2e17cf6af4974dc4f96f982c762fbd056b369608aece84f3c97eed6c02b7648d28a1ab3f3566026c830a6eb4d57a55676a251768961021c9050a7b6974e74c7efd13ca9324ac1a6768bca2daa0daf12b3c8ebc3ad02e5d5c935024048da77f0b6047842ebceaf42ec881dd23660f5f4caad61c04ea8736f4b8abd5c81197ebd10853a45e4944b85d318500c385596344922b43e114698bc05474aad582c&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:38:04.830072606+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/config.yaml?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=c0195bb166fc1589772f7ac8000fd444b8e06323ba6372e465965830037d0e3abe33eb458ac66e538461c44cd0aad54fae2cf3ed0f374b29b4cc12eefa1bd24dbcc6dd7346e5d2fe8238aca5f9f454f181d5839a406a34dc8e519e37ef8c4fc8de26b075401ceec8431d3f6c56345d424b5cace6503127369284bda44efb6fe17a73b78fc162b5452942692cd51024f02807e743229ac8d7ddee751592f125d012c3e4de0790abc7b061c8d8d6916a275902955479dcddb6342607f857379f4552e01426a19bf54b2231de5fb22f340330a26d23e6cae7d90831b2172a8cfdd004a5de1aca9d47db2eded16ebcbeb22d617adfb1a0324e51847993c2374681dd&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:43546->142.250.73.123:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/config.yaml?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=c0195bb166fc1589772f7ac8000fd444b8e06323ba6372e465965830037d0e3abe33eb458ac66e538461c44cd0aad54fae2cf3ed0f374b29b4cc12eefa1bd24dbcc6dd7346e5d2fe8238aca5f9f454f181d5839a406a34dc8e519e37ef8c4fc8de26b075401ceec8431d3f6c56345d424b5cace6503127369284bda44efb6fe17a73b78fc162b5452942692cd51024f02807e743229ac8d7ddee751592f125d012c3e4de0790abc7b061c8d8d6916a275902955479dcddb6342607f857379f4552e01426a19bf54b2231de5fb22f340330a26d23e6cae7d90831b2172a8cfdd004a5de1aca9d47db2eded16ebcbeb22d617adfb1a0324e51847993c2374681dd&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:38:08.914871723+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=1c9fbc6cbd51ae71160b439aa8a692f95ee16291fd45e2c33912f90a5ae3a56a09f33e745a416b04da501c1cf40827caa3399f6cde0752ce3f571c12353aa477d27f258d11b74762a029b2fada9797f27a85c7fe03e3bb1ec12fb98789a6b389cf000b2742ca658544ea41f22c9631272e67f2ea275c4a5b7b9c516fed6ed22ccce6b26b684c341e337963a684be9b248bd64b665fb73437d28ce7365767fe6721770c9f8f75e150767d9e4154a983b7f6aafcc31bbc050a26df2e9a3b99e065c111d7a92525568b04cf8d8a7490e6b8f77aed12c26aa378d0708385b0f40b22af3988b5c6cfb25918b1c9471be072bdfdc7a986ca0774f732d27361d23f3448&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:40682->142.250.73.91:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=1c9fbc6cbd51ae71160b439aa8a692f95ee16291fd45e2c33912f90a5ae3a56a09f33e745a416b04da501c1cf40827caa3399f6cde0752ce3f571c12353aa477d27f258d11b74762a029b2fada9797f27a85c7fe03e3bb1ec12fb98789a6b389cf000b2742ca658544ea41f22c9631272e67f2ea275c4a5b7b9c516fed6ed22ccce6b26b684c341e337963a684be9b248bd64b665fb73437d28ce7365767fe6721770c9f8f75e150767d9e4154a983b7f6aafcc31bbc050a26df2e9a3b99e065c111d7a92525568b04cf8d8a7490e6b8f77aed12c26aa378d0708385b0f40b22af3988b5c6cfb25918b1c9471be072bdfdc7a986ca0774f732d27361d23f3448&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:38:09.512888171+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/requirements.txt?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=27041959870a9e2ba8f974ca28bb5723ad374fc1917b982522ea3f555344b25fed8f73d5b443663eb1d7d10499d748efa90c8f6d91802daabf9e8c31258a7371e93aca1694d3e8bcb35a72d7c9ca243dc164f5dfae6ec3009c25cc78ea8cf37629f017b7538998f44f7c65ccfb675e343601421475a4490c754c5ee0370d3d5dfa928faddfe9a90621302ce69efd3d26c51f49c23f92148b018281ccd02f22c42e73e318594ea9c2ff9b25ad13163b60c37f1ededb4dc3a50712bcabdffed71883e2e5b5f04c40f13612d9ff1b6762b9f79d6e19873e959fbf9495e4827e901e8dbd25d7b6f291841915bc428212ad0a142cc3a9b7d04dbba9b49873573481fa&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:53738->142.250.73.123:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/requirements.txt?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=27041959870a9e2ba8f974ca28bb5723ad374fc1917b982522ea3f555344b25fed8f73d5b443663eb1d7d10499d748efa90c8f6d91802daabf9e8c31258a7371e93aca1694d3e8bcb35a72d7c9ca243dc164f5dfae6ec3009c25cc78ea8cf37629f017b7538998f44f7c65ccfb675e343601421475a4490c754c5ee0370d3d5dfa928faddfe9a90621302ce69efd3d26c51f49c23f92148b018281ccd02f22c42e73e318594ea9c2ff9b25ad13163b60c37f1ededb4dc3a50712bcabdffed71883e2e5b5f04c40f13612d9ff1b6762b9f79d6e19873e959fbf9495e4827e901e8dbd25d7b6f291841915bc428212ad0a142cc3a9b7d04dbba9b49873573481fa&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:38:25.342508911+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=1c9fbc6cbd51ae71160b439aa8a692f95ee16291fd45e2c33912f90a5ae3a56a09f33e745a416b04da501c1cf40827caa3399f6cde0752ce3f571c12353aa477d27f258d11b74762a029b2fada9797f27a85c7fe03e3bb1ec12fb98789a6b389cf000b2742ca658544ea41f22c9631272e67f2ea275c4a5b7b9c516fed6ed22ccce6b26b684c341e337963a684be9b248bd64b665fb73437d28ce7365767fe6721770c9f8f75e150767d9e4154a983b7f6aafcc31bbc050a26df2e9a3b99e065c111d7a92525568b04cf8d8a7490e6b8f77aed12c26aa378d0708385b0f40b22af3988b5c6cfb25918b1c9471be072bdfdc7a986ca0774f732d27361d23f3448&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:49916->142.250.73.91:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=1c9fbc6cbd51ae71160b439aa8a692f95ee16291fd45e2c33912f90a5ae3a56a09f33e745a416b04da501c1cf40827caa3399f6cde0752ce3f571c12353aa477d27f258d11b74762a029b2fada9797f27a85c7fe03e3bb1ec12fb98789a6b389cf000b2742ca658544ea41f22c9631272e67f2ea275c4a5b7b9c516fed6ed22ccce6b26b684c341e337963a684be9b248bd64b665fb73437d28ce7365767fe6721770c9f8f75e150767d9e4154a983b7f6aafcc31bbc050a26df2e9a3b99e065c111d7a92525568b04cf8d8a7490e6b8f77aed12c26aa378d0708385b0f40b22af3988b5c6cfb25918b1c9471be072bdfdc7a986ca0774f732d27361d23f3448&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:38:59.014612034+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=1c9fbc6cbd51ae71160b439aa8a692f95ee16291fd45e2c33912f90a5ae3a56a09f33e745a416b04da501c1cf40827caa3399f6cde0752ce3f571c12353aa477d27f258d11b74762a029b2fada9797f27a85c7fe03e3bb1ec12fb98789a6b389cf000b2742ca658544ea41f22c9631272e67f2ea275c4a5b7b9c516fed6ed22ccce6b26b684c341e337963a684be9b248bd64b665fb73437d28ce7365767fe6721770c9f8f75e150767d9e4154a983b7f6aafcc31bbc050a26df2e9a3b99e065c111d7a92525568b04cf8d8a7490e6b8f77aed12c26aa378d0708385b0f40b22af3988b5c6cfb25918b1c9471be072bdfdc7a986ca0774f732d27361d23f3448&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:44582->142.250.73.123:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073326Z&X-Goog-Expires=86399&X-Goog-Signature=1c9fbc6cbd51ae71160b439aa8a692f95ee16291fd45e2c33912f90a5ae3a56a09f33e745a416b04da501c1cf40827caa3399f6cde0752ce3f571c12353aa477d27f258d11b74762a029b2fada9797f27a85c7fe03e3bb1ec12fb98789a6b389cf000b2742ca658544ea41f22c9631272e67f2ea275c4a5b7b9c516fed6ed22ccce6b26b684c341e337963a684be9b248bd64b665fb73437d28ce7365767fe6721770c9f8f75e150767d9e4154a983b7f6aafcc31bbc050a26df2e9a3b99e065c111d7a92525568b04cf8d8a7490e6b8f77aed12c26aa378d0708385b0f40b22af3988b5c6cfb25918b1c9471be072bdfdc7a986ca0774f732d27361d23f3448&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:40:15.762761262+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/output.log?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=1d3cd94b4887f4c0adaa644487d39281fc878818fa20fa0c39b307b4d29a0774e1077acff64b5047465d575c9417f96645240e67e6dab6d1cc25f617e622bcf0f8715629cc8509684074bcd0688360af452c0f89dc80a4e9025500a6e0d16485ec7617abfd0ff10c8ae62cf6012c0f43aad8d30bb7c9de79f87c56f73ccabc572e83285debd09c51e5fcc95e1a7fe1cc011606044e056dc5f2d6756e365041194a3e52cc500d0e71499406481acf9cc760c1da018d07a35e07830dd04208781140f1fc38c86f5aa24bf65e1a58d35406ae84e7026fe6cf59353cf16ca53b51a813beb22417486e4fabc77de1aaef6ee20b755ce21b765571f6bfbf8473942f9d&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:46682->142.250.73.91:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/output.log?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=1d3cd94b4887f4c0adaa644487d39281fc878818fa20fa0c39b307b4d29a0774e1077acff64b5047465d575c9417f96645240e67e6dab6d1cc25f617e622bcf0f8715629cc8509684074bcd0688360af452c0f89dc80a4e9025500a6e0d16485ec7617abfd0ff10c8ae62cf6012c0f43aad8d30bb7c9de79f87c56f73ccabc572e83285debd09c51e5fcc95e1a7fe1cc011606044e056dc5f2d6756e365041194a3e52cc500d0e71499406481acf9cc760c1da018d07a35e07830dd04208781140f1fc38c86f5aa24bf65e1a58d35406ae84e7026fe6cf59353cf16ca53b51a813beb22417486e4fabc77de1aaef6ee20b755ce21b765571f6bfbf8473942f9d&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:40:15.7960472+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-summary.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=43378ba7835541b25e7884ddd1295c99cdb59af218b575e11ebf49219cd7b70414883b2eaabf6ce5bf29b6dcf3aa4854cc71ac3f00aceb098b2bd9f008cbf3c3e776aeb0eec09eafbd1633dbdf5cd30ca267e7fc78a69846132727fb83834b13e3aaf2e17cf6af4974dc4f96f982c762fbd056b369608aece84f3c97eed6c02b7648d28a1ab3f3566026c830a6eb4d57a55676a251768961021c9050a7b6974e74c7efd13ca9324ac1a6768bca2daa0daf12b3c8ebc3ad02e5d5c935024048da77f0b6047842ebceaf42ec881dd23660f5f4caad61c04ea8736f4b8abd5c81197ebd10853a45e4944b85d318500c385596344922b43e114698bc05474aad582c&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:46696->142.250.73.91:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/wandb-summary.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=43378ba7835541b25e7884ddd1295c99cdb59af218b575e11ebf49219cd7b70414883b2eaabf6ce5bf29b6dcf3aa4854cc71ac3f00aceb098b2bd9f008cbf3c3e776aeb0eec09eafbd1633dbdf5cd30ca267e7fc78a69846132727fb83834b13e3aaf2e17cf6af4974dc4f96f982c762fbd056b369608aece84f3c97eed6c02b7648d28a1ab3f3566026c830a6eb4d57a55676a251768961021c9050a7b6974e74c7efd13ca9324ac1a6768bca2daa0daf12b3c8ebc3ad02e5d5c935024048da77f0b6047842ebceaf42ec881dd23660f5f4caad61c04ea8736f4b8abd5c81197ebd10853a45e4944b85d318500c385596344922b43e114698bc05474aad582c&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:40:15.806561829+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/config.yaml?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=c0195bb166fc1589772f7ac8000fd444b8e06323ba6372e465965830037d0e3abe33eb458ac66e538461c44cd0aad54fae2cf3ed0f374b29b4cc12eefa1bd24dbcc6dd7346e5d2fe8238aca5f9f454f181d5839a406a34dc8e519e37ef8c4fc8de26b075401ceec8431d3f6c56345d424b5cace6503127369284bda44efb6fe17a73b78fc162b5452942692cd51024f02807e743229ac8d7ddee751592f125d012c3e4de0790abc7b061c8d8d6916a275902955479dcddb6342607f857379f4552e01426a19bf54b2231de5fb22f340330a26d23e6cae7d90831b2172a8cfdd004a5de1aca9d47db2eded16ebcbeb22d617adfb1a0324e51847993c2374681dd&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:46686->142.250.73.91:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/config.yaml?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=c0195bb166fc1589772f7ac8000fd444b8e06323ba6372e465965830037d0e3abe33eb458ac66e538461c44cd0aad54fae2cf3ed0f374b29b4cc12eefa1bd24dbcc6dd7346e5d2fe8238aca5f9f454f181d5839a406a34dc8e519e37ef8c4fc8de26b075401ceec8431d3f6c56345d424b5cace6503127369284bda44efb6fe17a73b78fc162b5452942692cd51024f02807e743229ac8d7ddee751592f125d012c3e4de0790abc7b061c8d8d6916a275902955479dcddb6342607f857379f4552e01426a19bf54b2231de5fb22f340330a26d23e6cae7d90831b2172a8cfdd004a5de1aca9d47db2eded16ebcbeb22d617adfb1a0324e51847993c2374681dd&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:40:20.882189315+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/config.yaml?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=c0195bb166fc1589772f7ac8000fd444b8e06323ba6372e465965830037d0e3abe33eb458ac66e538461c44cd0aad54fae2cf3ed0f374b29b4cc12eefa1bd24dbcc6dd7346e5d2fe8238aca5f9f454f181d5839a406a34dc8e519e37ef8c4fc8de26b075401ceec8431d3f6c56345d424b5cace6503127369284bda44efb6fe17a73b78fc162b5452942692cd51024f02807e743229ac8d7ddee751592f125d012c3e4de0790abc7b061c8d8d6916a275902955479dcddb6342607f857379f4552e01426a19bf54b2231de5fb22f340330a26d23e6cae7d90831b2172a8cfdd004a5de1aca9d47db2eded16ebcbeb22d617adfb1a0324e51847993c2374681dd&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.1.90:55506->142.250.217.91:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/ig4rhoqf/config.yaml?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250813%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250813T073342Z&X-Goog-Expires=86399&X-Goog-Signature=c0195bb166fc1589772f7ac8000fd444b8e06323ba6372e465965830037d0e3abe33eb458ac66e538461c44cd0aad54fae2cf3ed0f374b29b4cc12eefa1bd24dbcc6dd7346e5d2fe8238aca5f9f454f181d5839a406a34dc8e519e37ef8c4fc8de26b075401ceec8431d3f6c56345d424b5cace6503127369284bda44efb6fe17a73b78fc162b5452942692cd51024f02807e743229ac8d7ddee751592f125d012c3e4de0790abc7b061c8d8d6916a275902955479dcddb6342607f857379f4552e01426a19bf54b2231de5fb22f340330a26d23e6cae7d90831b2172a8cfdd004a5de1aca9d47db2eded16ebcbeb22d617adfb1a0324e51847993c2374681dd&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-13T15:42:32.350288869+08:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"}
+{"time":"2025-08-13T15:42:37.496789088+08:00","level":"INFO","msg":"handler: closed","stream_id":"ig4rhoqf"}
+{"time":"2025-08-13T15:42:37.502696559+08:00","level":"INFO","msg":"sender: closed","stream_id":"ig4rhoqf"}
+{"time":"2025-08-13T15:42:37.502726375+08:00","level":"INFO","msg":"stream: closed","id":"ig4rhoqf"}

BioReason_new/wandb/debug.log

@@ -0,0 +1,23 @@
+2025-08-13 15:33:18,716 INFO    MainThread:13510 [wandb_setup.py:_flush():80] Current SDK version is 0.21.1
+2025-08-13 15:33:18,716 INFO    MainThread:13510 [wandb_setup.py:_flush():80] Configure stats pid to 13510
+2025-08-13 15:33:18,716 INFO    MainThread:13510 [wandb_setup.py:_flush():80] Loading settings from /root/.config/wandb/settings
+2025-08-13 15:33:18,716 INFO    MainThread:13510 [wandb_setup.py:_flush():80] Loading settings from /nas/shared/kilab/wangyujia/BioReason_new/wandb/settings
+2025-08-13 15:33:18,716 INFO    MainThread:13510 [wandb_setup.py:_flush():80] Loading settings from environment variables
+2025-08-13 15:33:18,716 INFO    MainThread:13510 [wandb_init.py:setup_run_log_directory():703] Logging user logs to /nas/shared/kilab/wangyujia/BioReason_new/wandb/run-20250813_153318-ig4rhoqf/logs/debug.log
+2025-08-13 15:33:18,716 INFO    MainThread:13510 [wandb_init.py:setup_run_log_directory():704] Logging internal logs to /nas/shared/kilab/wangyujia/BioReason_new/wandb/run-20250813_153318-ig4rhoqf/logs/debug-internal.log
+2025-08-13 15:33:18,716 INFO    MainThread:13510 [wandb_init.py:init():830] calling init triggers
+2025-08-13 15:33:18,716 INFO    MainThread:13510 [wandb_init.py:init():835] wandb.init called with sweep_config: {}
+config: {'text_model_name': '/oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged', 'protein_model_name': '/nas/shared/kilab/wangyujia/ProtT3/plm_model/esm2-150m', 'qformer_model_name': '/nas/shared/kilab/wangyujia/ProtT3/plm_model/microsoft', 'cache_dir': '/model-weights', 'num_query_tokens': 8, 'train_dataset': '/nas/shared/kilab/wangyujia/ProtT3/data/SwissProtV3/train_set.jsonl', 'valid_dataset': '/nas/shared/kilab/wangyujia/ProtT3/data/SwissProtV3/valid_set.jsonl', 'eval_dataset': True, 'output_dir': './contrastive_outputs', 'num_epochs': 10, 'batch_size': 8, 'learning_rate': 0.0001, 'weight_decay': 0.01, 'warmup_steps': 1000, 'gradient_accumulation_steps': 1, 'temperature': 0.07, 'freeze_protein_model': True, 'freeze_text_model': True, 'protein_weight': 1.0, 'text_weight': 1.0, 'enable_ptm': True, 'ptm_weight': 1.0, 'max_length_protein': 1024, 'max_length_text': 512, 'num_workers': 8, 'logging_steps': 100, 'eval_steps': 500, 'save_steps': 1000, 'save_total_limit': 3, 'fp16': False, 'bf16': False, 'seed': 42, 'use_wandb': True, 'wandb_project': 'protein-llm-contrastive', 'wandb_entity': None, '_wandb': {}}
+2025-08-13 15:33:18,716 INFO    MainThread:13510 [wandb_init.py:init():871] starting backend
+2025-08-13 15:33:18,925 INFO    MainThread:13510 [wandb_init.py:init():874] sending inform_init request
+2025-08-13 15:33:18,931 INFO    MainThread:13510 [wandb_init.py:init():882] backend started and connected
+2025-08-13 15:33:18,933 INFO    MainThread:13510 [wandb_init.py:init():953] updated telemetry
+2025-08-13 15:33:18,961 INFO    MainThread:13510 [wandb_init.py:init():977] communicating run to backend with 90.0 second timeout
+2025-08-13 15:33:25,902 INFO    MainThread:13510 [wandb_init.py:init():1029] starting run threads in backend
+2025-08-13 15:33:26,011 INFO    MainThread:13510 [wandb_run.py:_console_start():2494] atexit reg
+2025-08-13 15:33:26,011 INFO    MainThread:13510 [wandb_run.py:_redirect():2342] redirect: wrap_raw
+2025-08-13 15:33:26,011 INFO    MainThread:13510 [wandb_run.py:_redirect():2411] Wrapping output streams.
+2025-08-13 15:33:26,011 INFO    MainThread:13510 [wandb_run.py:_redirect():2434] Redirects installed.
+2025-08-13 15:33:26,014 INFO    MainThread:13510 [wandb_init.py:init():1075] run started, returning control to user process
+2025-08-13 15:33:36,032 INFO    MainThread:13510 [wandb_run.py:_config_callback():1380] config_cb None None {'output_dir': './contrastive_outputs', 'overwrite_output_dir': False, 'do_train': False, 'do_eval': True, 'do_predict': False, 'eval_strategy': 'steps', 'prediction_loss_only': False, 'per_device_train_batch_size': 8, 'per_device_eval_batch_size': 8, 'per_gpu_train_batch_size': None, 'per_gpu_eval_batch_size': None, 'gradient_accumulation_steps': 1, 'eval_accumulation_steps': None, 'eval_delay': 0, 'torch_empty_cache_steps': None, 'learning_rate': 0.0001, 'weight_decay': 0.01, 'adam_beta1': 0.9, 'adam_beta2': 0.999, 'adam_epsilon': 1e-08, 'max_grad_norm': 1.0, 'num_train_epochs': 10, 'max_steps': -1, 'lr_scheduler_type': 'linear', 'lr_scheduler_kwargs': {}, 'warmup_ratio': 0.0, 'warmup_steps': 1000, 'log_level': 'passive', 'log_level_replica': 'warning', 'log_on_each_node': True, 'logging_dir': './contrastive_outputs/runs/Aug13_15-33-30_dsw-265304-585cc9d768-ckfd7', 'logging_strategy': 'steps', 'logging_first_step': False, 'logging_steps': 100, 'logging_nan_inf_filter': True, 'save_strategy': 'steps', 'save_steps': 1000, 'save_total_limit': 3, 'save_safetensors': True, 'save_on_each_node': False, 'save_only_model': False, 'restore_callback_states_from_checkpoint': False, 'no_cuda': False, 'use_cpu': False, 'use_mps_device': False, 'seed': 42, 'data_seed': None, 'jit_mode_eval': False, 'use_ipex': False, 'bf16': False, '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': 500, 'dataloader_num_workers': 8, 'dataloader_prefetch_factor': None, 'past_index': -1, 'run_name': None, 'disable_tqdm': False, 'remove_unused_columns': False, 'label_names': None, 'load_best_model_at_end': True, 'metric_for_best_model': 'eval_avg_recall_at_1', 'greater_is_better': True, '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}, 'fsdp_transformer_layer_cls_to_wrap': None, 'accelerator_config': {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}, 'deepspeed': None, 'label_smoothing_factor': 0.0, 'optim': 'adamw_torch', 'optim_args': None, 'adafactor': False, 'group_by_length': False, 'length_column_name': 'length', 'report_to': ['wandb'], 'ddp_find_unused_parameters': False, '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': 'every_save', 'hub_token': '<HUB_TOKEN>', 'hub_private_repo': None, 'hub_always_push': False, 'hub_revision': None, 'gradient_checkpointing': False, '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': '<PUSH_TO_HUB_TOKEN>', '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': False, 'include_num_input_tokens_seen': False, 'neftune_noise_alpha': None, 'optim_target_modules': None, 'batch_eval_metrics': False, 'eval_on_start': False, 'use_liger_kernel': False, 'liger_kernel_config': None, 'eval_use_gather_object': False, 'average_tokens_across_devices': False, 'temperature': 0.07, 'freeze_protein_model': True, 'freeze_text_model': True, 'protein_weight': 1.0, 'text_weight': 1.0, 'max_length_protein': 1024, 'max_length_text': 512, 'enable_ptm': True, 'ptm_weight': 1.0}
+2025-08-13 15:33:41,038 INFO    MsgRouterThr:13510 [mailbox.py:close():129] [no run ID] Closing mailbox, abandoning 2 handles.

BioReason_new/wandb/run-20250811_215805-k21eogb7/files/config.yaml

@@ -0,0 +1,159 @@
+_wandb:
+    value:
+        cli_version: 0.21.1
+        e:
+            ryu7ghs3jgr22n9qfads9dl5ddd2hfe3:
+                args:
+                    - --text_model_name
+                    - /oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged
+                    - --protein_model_name
+                    - /nas/shared/kilab/wangyujia/ProtT3/plm_model/esm2-150m
+                    - --qformer_model_name
+                    - /nas/shared/kilab/wangyujia/ProtT3/plm_model/microsoft
+                    - --num_query_tokens
+                    - "8"
+                    - --dataset_name
+                    - /nas/shared/kilab/wangyujia/ProtT3/data/SwissProtV3/train_set.jsonl
+                    - --output_dir
+                    - ./contrastive_outputs
+                    - --num_epochs
+                    - "10"
+                    - --batch_size
+                    - "32"
+                    - --learning_rate
+                    - "1e-4"
+                    - --temperature
+                    - "0.07"
+                    - --freeze_protein_model
+                    - --freeze_text_model
+                    - --enable_ptm
+                    - --max_length_protein
+                    - "1024"
+                    - --max_length_text
+                    - "512"
+                    - --num_workers
+                    - "8"
+                    - --eval_dataset
+                    - --use_wandb
+                    - --wandb_project
+                    - protein-llm-contrastive
+                    - --logging_steps
+                    - "100"
+                    - --eval_steps
+                    - "500"
+                    - --save_steps
+                    - "1000"
+                codePath: wangyujia/BioReason_new/train_contrastive.py
+                codePathLocal: train_contrastive.py
+                executable: /root/miniconda3/envs/bioreason/bin/python
+                git:
+                    commit: b8caf406aa1699c788f0ca6e44a1769452c317db
+                    remote: https://github.com/PorUna-byte/PAR.git
+                host: dsw-265304-f8bc5ff76-4mdt5
+                os: Linux-5.10.134-008.18.kangaroo.al8.x86_64-x86_64-with-glibc2.35
+                program: /nas/shared/kilab/wangyujia/BioReason_new/train_contrastive.py
+                python: CPython 3.11.0
+                root: /nas/shared/kilab/wangyujia/BioReason_new
+                startedAt: "2025-08-11T13:58:05.851565Z"
+                writerId: ryu7ghs3jgr22n9qfads9dl5ddd2hfe3
+        m: []
+        python_version: 3.11.0
+        t:
+            "1":
+                - 1
+                - 9
+                - 11
+                - 41
+                - 49
+                - 51
+                - 71
+                - 84
+                - 98
+                - 103
+            "2":
+                - 1
+                - 9
+                - 11
+                - 41
+                - 49
+                - 51
+                - 71
+                - 84
+                - 98
+                - 103
+            "3":
+                - 13
+                - 16
+            "4": 3.11.0
+            "5": 0.21.1
+            "6": 4.55.0
+            "12": 0.21.1
+            "13": linux-x86_64
+batch_size:
+    value: 32
+bf16:
+    value: false
+cache_dir:
+    value: /model-weights
+dataset_name:
+    value: /nas/shared/kilab/wangyujia/ProtT3/data/SwissProtV3/train_set.jsonl
+enable_ptm:
+    value: true
+eval_dataset:
+    value: true
+eval_steps:
+    value: 500
+fp16:
+    value: false
+freeze_protein_model:
+    value: true
+freeze_text_model:
+    value: true
+gradient_accumulation_steps:
+    value: 1
+learning_rate:
+    value: 0.0001
+logging_steps:
+    value: 100
+max_length_protein:
+    value: 1024
+max_length_text:
+    value: 512
+num_epochs:
+    value: 10
+num_query_tokens:
+    value: 8
+num_workers:
+    value: 8
+output_dir:
+    value: ./contrastive_outputs
+protein_model_name:
+    value: /nas/shared/kilab/wangyujia/ProtT3/plm_model/esm2-150m
+protein_weight:
+    value: 1
+ptm_weight:
+    value: 1
+qformer_model_name:
+    value: /nas/shared/kilab/wangyujia/ProtT3/plm_model/microsoft
+save_steps:
+    value: 1000
+save_total_limit:
+    value: 3
+seed:
+    value: 42
+temperature:
+    value: 0.07
+text_model_name:
+    value: /oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged
+text_weight:
+    value: 1
+use_wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: protein-llm-contrastive
+warmup_steps:
+    value: 1000
+weight_decay:
+    value: 0.01

BioReason_new/wandb/run-20250811_215805-k21eogb7/files/output.log

@@ -0,0 +1,21 @@
+Loading model...
+Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████| 4/4 [00:01<00:00,  2.30it/s]
+Some weights of EsmModel were not initialized from the model checkpoint at /nas/shared/kilab/wangyujia/ProtT3/plm_model/esm2-150m and are newly initialized: ['pooler.dense.bias', 'pooler.dense.weight']
+You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
+Loading datasets...
+Traceback (most recent call last):
+  File "/nas/shared/kilab/wangyujia/BioReason_new/train_contrastive.py", line 549, in <module>
+    trainer = main(args)
+              ^^^^^^^^^^
+  File "/nas/shared/kilab/wangyujia/BioReason_new/train_contrastive.py", line 317, in main
+    train_dataset = load_dataset(args.dataset_name, split="train")
+                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/root/miniconda3/envs/bioreason/lib/python3.11/site-packages/datasets/load.py", line 1392, in load_dataset
+    builder_instance = load_dataset_builder(
+                       ^^^^^^^^^^^^^^^^^^^^^
+  File "/root/miniconda3/envs/bioreason/lib/python3.11/site-packages/datasets/load.py", line 1132, in load_dataset_builder
+    dataset_module = dataset_module_factory(
+                     ^^^^^^^^^^^^^^^^^^^^^^^
+  File "/root/miniconda3/envs/bioreason/lib/python3.11/site-packages/datasets/load.py", line 1033, in dataset_module_factory
+    raise FileNotFoundError(f"Couldn't find any data file at {relative_to_absolute_path(path)}.")
+FileNotFoundError: Couldn't find any data file at /nas/shared/kilab/wangyujia/ProtT3/data/SwissProtV3/train_set.jsonl.

BioReason_new/wandb/run-20250811_215805-k21eogb7/files/requirements.txt

@@ -0,0 +1,233 @@
+nvidia-nccl-cu12==2.26.2
+cbor2==5.6.5
+jupyter_server==2.16.0
+nvidia-curand-cu12==10.3.7.77
+bleach==6.2.0
+py-cpuinfo==9.0.0
+llvmlite==0.44.0
+fsspec==2025.3.0
+uvloop==0.21.0
+rfc3986-validator==0.1.1
+smmap==5.0.2
+pip==25.1
+compressed-tensors==0.10.2
+ipython_pygments_lexers==1.1.1
+fastapi-cli==0.0.8
+filelock==3.18.0
+msgspec==0.19.0
+hjson==3.1.0
+markdown-it-py==3.0.0
+pyzmq==27.0.1
+interegular==0.3.3
+widgetsnbextension==4.0.14
+vllm==0.10.0
+ipykernel==6.30.1
+pydantic==2.11.7
+click==8.2.1
+torchvision==0.22.1
+fastapi-cloud-cli==0.1.5
+httpcore==1.0.9
+nvidia-cuda-nvrtc-cu12==12.6.77
+mdurl==0.1.2
+rich-toolkit==0.15.0
+Pygments==2.19.2
+pure_eval==0.2.3
+types-python-dateutil==2.9.0.20250809
+referencing==0.36.2
+jupyterlab_widgets==3.0.15
+typing-inspection==0.4.1
+stack-data==0.6.3
+jupyter_client==8.6.3
+regex==2025.7.33
+platformdirs==4.3.8
+arrow==1.3.0
+aiosignal==1.4.0
+python-dateutil==2.9.0.post0
+numpy==2.2.6
+jupyter-lsp==2.2.6
+transformers==4.55.0
+mpmath==1.3.0
+six==1.17.0
+python-json-logger==3.3.0
+distro==1.9.0
+partial-json-parser==0.2.1.1.post6
+bitsandbytes==0.46.1
+nvidia-cusparselt-cu12==0.6.3
+pandocfilters==1.5.1
+pexpect==4.9.0
+pydantic-extra-types==2.10.5
+Jinja2==3.1.6
+sentencepiece==0.2.0
+uvicorn==0.35.0
+babel==2.17.0
+trl==0.21.0
+urllib3==2.5.0
+prometheus_client==0.22.1
+watchfiles==1.1.0
+prometheus-fastapi-instrumentator==7.1.0
+jsonschema-specifications==2025.4.1
+diskcache==5.6.3
+webcolors==24.11.1
+peft==0.17.0
+jiter==0.10.0
+triton==3.3.1
+gitdb==4.0.12
+gguf==0.17.1
+safetensors==0.6.2
+cloudpickle==3.1.1
+multiprocess==0.70.16
+aiohttp==3.12.15
+tornado==6.5.2
+nvidia-nvtx-cu12==12.6.77
+nbclient==0.10.2
+nbconvert==7.16.6
+psutil==7.0.0
+llguidance==0.7.30
+ray==2.48.0
+wcwidth==0.2.13
+rignore==0.6.4
+nvidia-cudnn-cu12==9.5.1.17
+soupsieve==2.7
+wandb==0.21.1
+overrides==7.7.0
+opencv-python-headless==4.12.0.88
+pycparser==2.22
+scipy==1.16.1
+terminado==0.18.1
+typer==0.16.0
+parso==0.8.4
+lark==1.2.2
+msgpack==1.1.1
+websockets==15.0.1
+idna==3.10
+fastrlock==0.8.3
+jedi==0.19.2
+accelerate==1.10.0
+jupyter==1.1.1
+beautifulsoup4==4.13.4
+h11==0.16.0
+MarkupSafe==3.0.2
+python-dotenv==1.1.1
+aiohappyeyeballs==2.6.1
+rich==14.1.0
+nbformat==5.10.4
+traitlets==5.14.3
+decorator==5.2.1
+soxr==0.5.0.post1
+propcache==0.3.2
+ninja==1.11.1.4
+cffi==1.17.1
+cupy-cuda12x==13.5.1
+pandas==2.3.1
+deepspeed==0.17.4
+setuptools==78.1.1
+websocket-client==1.8.0
+qwen-vl-utils==0.0.11
+webencodings==0.5.1
+httptools==0.6.4
+jupyterlab==4.4.5
+ptyprocess==0.7.0
+shellingham==1.5.4
+attrs==25.3.0
+fqdn==1.5.1
+huggingface-hub==0.34.4
+tokenizers==0.21.4
+asttokens==3.0.0
+jupyter_server_terminals==0.5.3
+av==15.0.0
+nvidia-cuda-cupti-cu12==12.6.80
+typing_extensions==4.14.1
+hf-xet==1.1.7
+jupyter_core==5.8.1
+starlette==0.47.2
+fastjsonschema==2.21.1
+fastapi==0.116.1
+lightning-utilities==0.15.2
+jupyter-console==6.6.3
+pybase64==1.4.2
+jupyter-events==0.12.0
+requests==2.32.4
+numba==0.61.2
+networkx==3.5
+nvidia-cusparse-cu12==12.5.4.2
+jsonpointer==3.0.0
+pyarrow==21.0.0
+dnspython==2.7.0
+torchaudio==2.7.1
+ipython==9.4.0
+isoduration==20.11.0
+bioreason==0.1.0
+matplotlib-inline==0.1.7
+packaging==25.0
+xxhash==3.5.0
+depyf==0.19.0
+sentry-sdk==2.34.1
+prompt_toolkit==3.0.51
+nvidia-cublas-cu12==12.6.4.1
+rfc3339-validator==0.1.4
+nvidia-cufft-cu12==11.3.0.4
+email_validator==2.2.0
+pycountry==24.6.1
+argon2-cffi==25.1.0
+nvidia-cufile-cu12==1.11.1.6
+frozenlist==1.7.0
+json5==0.12.0
+tinycss2==1.4.0
+defusedxml==0.7.1
+lm-format-enforcer==0.10.12
+Send2Trash==1.8.3
+anyio==4.10.0
+rfc3987-syntax==1.1.0
+pydantic_core==2.33.2
+debugpy==1.8.16
+async-lru==2.0.5
+nvidia-cuda-runtime-cu12==12.6.77
+tiktoken==0.11.0
+comm==0.2.3
+PyYAML==6.0.2
+blake3==1.0.5
+nvidia-cusolver-cu12==11.7.1.2
+torch==2.7.1
+torchmetrics==1.8.1
+yarl==1.20.1
+dill==0.3.8
+wheel==0.45.1
+cachetools==6.1.0
+multidict==6.6.3
+pytz==2025.2
+pillow==11.3.0
+annotated-types==0.7.0
+astor==0.8.1
+nest-asyncio==1.6.0
+httpx==0.28.1
+argon2-cffi-bindings==25.1.0
+notebook_shim==0.2.4
+jsonschema==4.25.0
+python-multipart==0.0.20
+charset-normalizer==3.4.3
+tqdm==4.67.1
+xformers==0.0.31
+tzdata==2025.2
+einops==0.8.1
+mistral_common==1.8.3
+jupyterlab_server==2.27.3
+sympy==1.14.0
+datasets==4.0.0
+GitPython==3.1.45
+mistune==3.1.3
+ipywidgets==8.1.7
+nvidia-ml-py==13.580.65
+uri-template==1.3.0
+notebook==7.4.5
+certifi==2025.8.3
+nvidia-nvjitlink-cu12==12.6.85
+openai==1.90.0
+xgrammar==0.1.21
+executing==2.2.0
+soundfile==0.13.1
+jupyterlab_pygments==0.3.0
+outlines_core==0.2.10
+sniffio==1.3.1
+pytorch-lightning==2.5.2
+rpds-py==0.27.0
+protobuf==6.31.1

BioReason_new/wandb/run-20250811_215805-k21eogb7/files/wandb-metadata.json

@@ -0,0 +1,57 @@
+{
+  "os": "Linux-5.10.134-008.18.kangaroo.al8.x86_64-x86_64-with-glibc2.35",
+  "python": "CPython 3.11.0",
+  "startedAt": "2025-08-11T13:58:05.851565Z",
+  "args": [
+    "--text_model_name",
+    "/oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged",
+    "--protein_model_name",
+    "/nas/shared/kilab/wangyujia/ProtT3/plm_model/esm2-150m",
+    "--qformer_model_name",
+    "/nas/shared/kilab/wangyujia/ProtT3/plm_model/microsoft",
+    "--num_query_tokens",
+    "8",
+    "--dataset_name",
+    "/nas/shared/kilab/wangyujia/ProtT3/data/SwissProtV3/train_set.jsonl",
+    "--output_dir",
+    "./contrastive_outputs",
+    "--num_epochs",
+    "10",
+    "--batch_size",
+    "32",
+    "--learning_rate",
+    "1e-4",
+    "--temperature",
+    "0.07",
+    "--freeze_protein_model",
+    "--freeze_text_model",
+    "--enable_ptm",
+    "--max_length_protein",
+    "1024",
+    "--max_length_text",
+    "512",
+    "--num_workers",
+    "8",
+    "--eval_dataset",
+    "--use_wandb",
+    "--wandb_project",
+    "protein-llm-contrastive",
+    "--logging_steps",
+    "100",
+    "--eval_steps",
+    "500",
+    "--save_steps",
+    "1000"
+  ],
+  "program": "/nas/shared/kilab/wangyujia/BioReason_new/train_contrastive.py",
+  "codePath": "wangyujia/BioReason_new/train_contrastive.py",
+  "codePathLocal": "train_contrastive.py",
+  "git": {
+    "remote": "https://github.com/PorUna-byte/PAR.git",
+    "commit": "b8caf406aa1699c788f0ca6e44a1769452c317db"
+  },
+  "root": "/nas/shared/kilab/wangyujia/BioReason_new",
+  "host": "dsw-265304-f8bc5ff76-4mdt5",
+  "executable": "/root/miniconda3/envs/bioreason/bin/python",
+  "writerId": "ryu7ghs3jgr22n9qfads9dl5ddd2hfe3"
+}

BioReason_new/wandb/run-20250811_215805-k21eogb7/files/wandb-summary.json

@@ -0,0 +1 @@
+{"_runtime":3,"_wandb":{"runtime":3}}

BioReason_new/wandb/run-20250811_215805-k21eogb7/logs/debug-internal.log

@@ -0,0 +1,15 @@
+{"time":"2025-08-11T21:58:06.266313778+08:00","level":"INFO","msg":"stream: starting","core version":"0.21.1"}
+{"time":"2025-08-11T21:58:11.67768592+08:00","level":"INFO","msg":"stream: created new stream","id":"k21eogb7"}
+{"time":"2025-08-11T21:58:11.67881+08:00","level":"INFO","msg":"stream: started","id":"k21eogb7"}
+{"time":"2025-08-11T21:58:11.678831626+08:00","level":"INFO","msg":"writer: started","stream_id":"k21eogb7"}
+{"time":"2025-08-11T21:58:11.678836167+08:00","level":"INFO","msg":"sender: started","stream_id":"k21eogb7"}
+{"time":"2025-08-11T21:58:11.678866578+08:00","level":"INFO","msg":"handler: started","stream_id":"k21eogb7"}
+{"time":"2025-08-11T21:58:17.422177077+08:00","level":"INFO","msg":"stream: closing","id":"k21eogb7"}
+{"time":"2025-08-11T21:58:17.894015835+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/k21eogb7/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250811%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250811T135814Z&X-Goog-Expires=86399&X-Goog-Signature=ac3cb37c5d4b2fc5862a0022639cf77d48b1d5bedf77614e65d21af8685129bd5888f0c2a5a385e3890e6bfe21ebfae49170fa950c85170f44ff15f9e914184a3755ccca89484e99e1c879cf0fe2606c7e0e29a42d8203a5625e24b124c0c3b9149353e88a5877f1c69381bbacde2febade041b46fa50ba30d73fd36d4edd17f8dbbc3e6d5589bf6a016fb4d5a7b8c061530322df6e2f1d2ce72f94f80d90e592b7ff11dee46dda00343f8c5cf527bfb758d250201795f13740771c5083590019038f3b4af4250b46cce19aa2c7d05c26aa69e1e6091738740ba50a56247d537ee611862178b0ecae0c476bb79d1e77eb7d472406c0ec7371a76080c28bdf798&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.8.118:37628->142.250.73.123:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/k21eogb7/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250811%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250811T135814Z&X-Goog-Expires=86399&X-Goog-Signature=ac3cb37c5d4b2fc5862a0022639cf77d48b1d5bedf77614e65d21af8685129bd5888f0c2a5a385e3890e6bfe21ebfae49170fa950c85170f44ff15f9e914184a3755ccca89484e99e1c879cf0fe2606c7e0e29a42d8203a5625e24b124c0c3b9149353e88a5877f1c69381bbacde2febade041b46fa50ba30d73fd36d4edd17f8dbbc3e6d5589bf6a016fb4d5a7b8c061530322df6e2f1d2ce72f94f80d90e592b7ff11dee46dda00343f8c5cf527bfb758d250201795f13740771c5083590019038f3b4af4250b46cce19aa2c7d05c26aa69e1e6091738740ba50a56247d537ee611862178b0ecae0c476bb79d1e77eb7d472406c0ec7371a76080c28bdf798&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-11T21:58:17.907861315+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/k21eogb7/requirements.txt?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250811%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250811T135814Z&X-Goog-Expires=86399&X-Goog-Signature=ba6b1293deff25490307723c6c45ae5a2004c6679aca654218f25fa93a552e701ef0117ff99cef40a8a866d8bfe1a7a30662ae892d1831dd9af3e37dbabb7395c8e38d0701570166647ba2ab06044cc1206ce74aac46ce4e3c16fe6ede3e068e29538f5ce545aa06ccc122ff37f2e0f230c5e9bb4c5f9e8fb37cb62a240476326613b447cf162d0a360167f2a2000f4ce82833edc6ed6fbea897f4f74b4210dcc3be2c47f92f8af94d520341231457a0a5796e6169943c0a3ac34cc0c5fc29141610ab6bb31c50000377c7bd4e21f805feb05e31212a54ca8a02be4a129bcd603dc585268f4e8ab0c5cadb4100d577cc3d6e9f38b36192b2f7222e2fe99f654a&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.8.118:37620->142.250.73.123:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/k21eogb7/requirements.txt?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250811%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250811T135814Z&X-Goog-Expires=86399&X-Goog-Signature=ba6b1293deff25490307723c6c45ae5a2004c6679aca654218f25fa93a552e701ef0117ff99cef40a8a866d8bfe1a7a30662ae892d1831dd9af3e37dbabb7395c8e38d0701570166647ba2ab06044cc1206ce74aac46ce4e3c16fe6ede3e068e29538f5ce545aa06ccc122ff37f2e0f230c5e9bb4c5f9e8fb37cb62a240476326613b447cf162d0a360167f2a2000f4ce82833edc6ed6fbea897f4f74b4210dcc3be2c47f92f8af94d520341231457a0a5796e6169943c0a3ac34cc0c5fc29141610ab6bb31c50000377c7bd4e21f805feb05e31212a54ca8a02be4a129bcd603dc585268f4e8ab0c5cadb4100d577cc3d6e9f38b36192b2f7222e2fe99f654a&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-11T21:58:20.579200782+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/k21eogb7/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250811%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250811T135814Z&X-Goog-Expires=86399&X-Goog-Signature=ac3cb37c5d4b2fc5862a0022639cf77d48b1d5bedf77614e65d21af8685129bd5888f0c2a5a385e3890e6bfe21ebfae49170fa950c85170f44ff15f9e914184a3755ccca89484e99e1c879cf0fe2606c7e0e29a42d8203a5625e24b124c0c3b9149353e88a5877f1c69381bbacde2febade041b46fa50ba30d73fd36d4edd17f8dbbc3e6d5589bf6a016fb4d5a7b8c061530322df6e2f1d2ce72f94f80d90e592b7ff11dee46dda00343f8c5cf527bfb758d250201795f13740771c5083590019038f3b4af4250b46cce19aa2c7d05c26aa69e1e6091738740ba50a56247d537ee611862178b0ecae0c476bb79d1e77eb7d472406c0ec7371a76080c28bdf798&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.8.118:51916->142.250.73.91:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/k21eogb7/wandb-metadata.json?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250811%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250811T135814Z&X-Goog-Expires=86399&X-Goog-Signature=ac3cb37c5d4b2fc5862a0022639cf77d48b1d5bedf77614e65d21af8685129bd5888f0c2a5a385e3890e6bfe21ebfae49170fa950c85170f44ff15f9e914184a3755ccca89484e99e1c879cf0fe2606c7e0e29a42d8203a5625e24b124c0c3b9149353e88a5877f1c69381bbacde2febade041b46fa50ba30d73fd36d4edd17f8dbbc3e6d5589bf6a016fb4d5a7b8c061530322df6e2f1d2ce72f94f80d90e592b7ff11dee46dda00343f8c5cf527bfb758d250201795f13740771c5083590019038f3b4af4250b46cce19aa2c7d05c26aa69e1e6091738740ba50a56247d537ee611862178b0ecae0c476bb79d1e77eb7d472406c0ec7371a76080c28bdf798&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-11T21:58:20.615403659+08:00","level":"ERROR","msg":"request failed","error":"Put \"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/k21eogb7/requirements.txt?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250811%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250811T135814Z&X-Goog-Expires=86399&X-Goog-Signature=ba6b1293deff25490307723c6c45ae5a2004c6679aca654218f25fa93a552e701ef0117ff99cef40a8a866d8bfe1a7a30662ae892d1831dd9af3e37dbabb7395c8e38d0701570166647ba2ab06044cc1206ce74aac46ce4e3c16fe6ede3e068e29538f5ce545aa06ccc122ff37f2e0f230c5e9bb4c5f9e8fb37cb62a240476326613b447cf162d0a360167f2a2000f4ce82833edc6ed6fbea897f4f74b4210dcc3be2c47f92f8af94d520341231457a0a5796e6169943c0a3ac34cc0c5fc29141610ab6bb31c50000377c7bd4e21f805feb05e31212a54ca8a02be4a129bcd603dc585268f4e8ab0c5cadb4100d577cc3d6e9f38b36192b2f7222e2fe99f654a&X-Goog-SignedHeaders=host&X-User=gia0603yucca\": read tcp 10.1.8.118:51920->142.250.73.91:443: read: connection reset by peer","method":"PUT","url":"https://storage.googleapis.com/wandb-production.appspot.com/gia0603yucca/protein-llm-contrastive/k21eogb7/requirements.txt?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=gorilla-files-url-signer-man%40wandb-production.iam.gserviceaccount.com%2F20250811%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20250811T135814Z&X-Goog-Expires=86399&X-Goog-Signature=ba6b1293deff25490307723c6c45ae5a2004c6679aca654218f25fa93a552e701ef0117ff99cef40a8a866d8bfe1a7a30662ae892d1831dd9af3e37dbabb7395c8e38d0701570166647ba2ab06044cc1206ce74aac46ce4e3c16fe6ede3e068e29538f5ce545aa06ccc122ff37f2e0f230c5e9bb4c5f9e8fb37cb62a240476326613b447cf162d0a360167f2a2000f4ce82833edc6ed6fbea897f4f74b4210dcc3be2c47f92f8af94d520341231457a0a5796e6169943c0a3ac34cc0c5fc29141610ab6bb31c50000377c7bd4e21f805feb05e31212a54ca8a02be4a129bcd603dc585268f4e8ab0c5cadb4100d577cc3d6e9f38b36192b2f7222e2fe99f654a&X-Goog-SignedHeaders=host&X-User=gia0603yucca"}
+{"time":"2025-08-11T21:58:35.386126561+08:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"}
+{"time":"2025-08-11T21:59:02.436399111+08:00","level":"INFO","msg":"handler: closed","stream_id":"k21eogb7"}
+{"time":"2025-08-11T21:59:02.440170612+08:00","level":"INFO","msg":"sender: closed","stream_id":"k21eogb7"}
+{"time":"2025-08-11T21:59:02.440183082+08:00","level":"INFO","msg":"stream: closed","id":"k21eogb7"}

BioReason_new/wandb/run-20250811_215805-k21eogb7/logs/debug.log

@@ -0,0 +1,22 @@
+2025-08-11 21:58:06,022 INFO    MainThread:79345 [wandb_setup.py:_flush():80] Current SDK version is 0.21.1
+2025-08-11 21:58:06,022 INFO    MainThread:79345 [wandb_setup.py:_flush():80] Configure stats pid to 79345
+2025-08-11 21:58:06,022 INFO    MainThread:79345 [wandb_setup.py:_flush():80] Loading settings from /root/.config/wandb/settings
+2025-08-11 21:58:06,022 INFO    MainThread:79345 [wandb_setup.py:_flush():80] Loading settings from /nas/shared/kilab/wangyujia/BioReason_new/wandb/settings
+2025-08-11 21:58:06,022 INFO    MainThread:79345 [wandb_setup.py:_flush():80] Loading settings from environment variables
+2025-08-11 21:58:06,022 INFO    MainThread:79345 [wandb_init.py:setup_run_log_directory():703] Logging user logs to /nas/shared/kilab/wangyujia/BioReason_new/wandb/run-20250811_215805-k21eogb7/logs/debug.log
+2025-08-11 21:58:06,022 INFO    MainThread:79345 [wandb_init.py:setup_run_log_directory():704] Logging internal logs to /nas/shared/kilab/wangyujia/BioReason_new/wandb/run-20250811_215805-k21eogb7/logs/debug-internal.log
+2025-08-11 21:58:06,022 INFO    MainThread:79345 [wandb_init.py:init():830] calling init triggers
+2025-08-11 21:58:06,022 INFO    MainThread:79345 [wandb_init.py:init():835] wandb.init called with sweep_config: {}
+config: {'text_model_name': '/oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged', 'protein_model_name': '/nas/shared/kilab/wangyujia/ProtT3/plm_model/esm2-150m', 'qformer_model_name': '/nas/shared/kilab/wangyujia/ProtT3/plm_model/microsoft', 'cache_dir': '/model-weights', 'num_query_tokens': 8, 'dataset_name': '/nas/shared/kilab/wangyujia/ProtT3/data/SwissProtV3/train_set.jsonl', 'eval_dataset': True, 'output_dir': './contrastive_outputs', 'num_epochs': 10, 'batch_size': 32, 'learning_rate': 0.0001, 'weight_decay': 0.01, 'warmup_steps': 1000, 'gradient_accumulation_steps': 1, 'temperature': 0.07, 'freeze_protein_model': True, 'freeze_text_model': True, 'protein_weight': 1.0, 'text_weight': 1.0, 'enable_ptm': True, 'ptm_weight': 1.0, 'max_length_protein': 1024, 'max_length_text': 512, 'num_workers': 8, 'logging_steps': 100, 'eval_steps': 500, 'save_steps': 1000, 'save_total_limit': 3, 'fp16': False, 'bf16': False, 'seed': 42, 'use_wandb': True, 'wandb_project': 'protein-llm-contrastive', 'wandb_entity': None, '_wandb': {}}
+2025-08-11 21:58:06,022 INFO    MainThread:79345 [wandb_init.py:init():871] starting backend
+2025-08-11 21:58:06,233 INFO    MainThread:79345 [wandb_init.py:init():874] sending inform_init request
+2025-08-11 21:58:06,259 INFO    MainThread:79345 [wandb_init.py:init():882] backend started and connected
+2025-08-11 21:58:06,263 INFO    MainThread:79345 [wandb_init.py:init():953] updated telemetry
+2025-08-11 21:58:06,327 INFO    MainThread:79345 [wandb_init.py:init():977] communicating run to backend with 90.0 second timeout
+2025-08-11 21:58:13,781 INFO    MainThread:79345 [wandb_init.py:init():1029] starting run threads in backend
+2025-08-11 21:58:13,892 INFO    MainThread:79345 [wandb_run.py:_console_start():2494] atexit reg
+2025-08-11 21:58:13,892 INFO    MainThread:79345 [wandb_run.py:_redirect():2342] redirect: wrap_raw
+2025-08-11 21:58:13,892 INFO    MainThread:79345 [wandb_run.py:_redirect():2411] Wrapping output streams.
+2025-08-11 21:58:13,892 INFO    MainThread:79345 [wandb_run.py:_redirect():2434] Redirects installed.
+2025-08-11 21:58:13,895 INFO    MainThread:79345 [wandb_init.py:init():1075] run started, returning control to user process
+2025-08-11 21:58:17,421 INFO    MsgRouterThr:79345 [mailbox.py:close():129] [no run ID] Closing mailbox, abandoning 2 handles.

BioReason_new/wandb/run-20250811_220309-2qgjwsxa/files/config.yaml

@@ -0,0 +1,195 @@
+_wandb:
+    value:
+        cli_version: 0.21.1
+        e:
+            r185oiuz6xjarzg7yyfap3b9flv6ll88:
+                args:
+                    - --text_model_name
+                    - /oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged
+                    - --protein_model_name
+                    - /nas/shared/kilab/wangyujia/ProtT3/plm_model/esm2-150m
+                    - --qformer_model_name
+                    - /nas/shared/kilab/wangyujia/ProtT3/plm_model/microsoft
+                    - --num_query_tokens
+                    - "8"
+                    - --dataset_name
+                    - /nas/shared/kilab/wangyujia/ProtT3/data/SwissProtV3/train_set.jsonl
+                    - --output_dir
+                    - ./contrastive_outputs
+                    - --num_epochs
+                    - "10"
+                    - --batch_size
+                    - "32"
+                    - --learning_rate
+                    - "1e-4"
+                    - --temperature
+                    - "0.07"
+                    - --freeze_protein_model
+                    - --freeze_text_model
+                    - --enable_ptm
+                    - --max_length_protein
+                    - "1024"
+                    - --max_length_text
+                    - "512"
+                    - --num_workers
+                    - "8"
+                    - --eval_dataset
+                    - --use_wandb
+                    - --wandb_project
+                    - protein-llm-contrastive
+                    - --logging_steps
+                    - "100"
+                    - --eval_steps
+                    - "500"
+                    - --save_steps
+                    - "1000"
+                codePath: wangyujia/BioReason_new/train_contrastive.py
+                codePathLocal: train_contrastive.py
+                cpu_count: 64
+                cpu_count_logical: 64
+                cudaVersion: "12.1"
+                disk:
+                    /:
+                        total: "1623302262784"
+                        used: "28193923072"
+                executable: /root/miniconda3/envs/bioreason/bin/python
+                git:
+                    commit: b8caf406aa1699c788f0ca6e44a1769452c317db
+                    remote: https://github.com/PorUna-byte/PAR.git
+                gpu: NVIDIA A800-SXM4-80GB
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ampere
+                      name: NVIDIA A800-SXM4-80GB
+                      uuid: GPU-71607f78-ad31-1ea4-19c1-908e3e31aaf1
+                    - architecture: Ampere
+                      name: NVIDIA A800-SXM4-80GB
+                      uuid: GPU-92b7dbbd-7ef5-3c5f-ce1c-1d179d7fa587
+                    - architecture: Ampere
+                      name: NVIDIA A800-SXM4-80GB
+                      uuid: GPU-bbc35439-ad79-578b-381b-aba6f0cc0168
+                    - architecture: Ampere
+                      name: NVIDIA A800-SXM4-80GB
+                      uuid: GPU-e492e147-ca2e-76f2-85da-4e08e4deeb14
+                    - architecture: Ampere
+                      name: NVIDIA A800-SXM4-80GB
+                      uuid: GPU-8c4f8e67-4b52-5107-3095-0f007e6378ac
+                    - architecture: Ampere
+                      name: NVIDIA A800-SXM4-80GB
+                      uuid: GPU-7063f0b9-4ca2-6a72-522d-1262899ac5ad
+                    - architecture: Ampere
+                      name: NVIDIA A800-SXM4-80GB
+                      uuid: GPU-3b6e9a37-bcf3-387c-7874-4f8de4abd115
+                    - architecture: Ampere
+                      name: NVIDIA A800-SXM4-80GB
+                      uuid: GPU-92456839-e814-7be9-6817-f3e8da8aa80c
+                host: dsw-265304-f8bc5ff76-4mdt5
+                memory:
+                    total: "549755813888"
+                os: Linux-5.10.134-008.18.kangaroo.al8.x86_64-x86_64-with-glibc2.35
+                program: /nas/shared/kilab/wangyujia/BioReason_new/train_contrastive.py
+                python: CPython 3.11.0
+                root: /nas/shared/kilab/wangyujia/BioReason_new
+                startedAt: "2025-08-11T14:03:09.687288Z"
+                writerId: r185oiuz6xjarzg7yyfap3b9flv6ll88
+        m: []
+        python_version: 3.11.0
+        t:
+            "1":
+                - 1
+                - 9
+                - 11
+                - 41
+                - 49
+                - 51
+                - 71
+                - 84
+                - 98
+                - 103
+            "2":
+                - 1
+                - 9
+                - 11
+                - 41
+                - 49
+                - 51
+                - 71
+                - 84
+                - 98
+                - 103
+            "3":
+                - 13
+                - 16
+            "4": 3.11.0
+            "5": 0.21.1
+            "6": 4.55.0
+            "12": 0.21.1
+            "13": linux-x86_64
+batch_size:
+    value: 32
+bf16:
+    value: false
+cache_dir:
+    value: /model-weights
+dataset_name:
+    value: /nas/shared/kilab/wangyujia/ProtT3/data/SwissProtV3/train_set.jsonl
+enable_ptm:
+    value: true
+eval_dataset:
+    value: true
+eval_steps:
+    value: 500
+fp16:
+    value: false
+freeze_protein_model:
+    value: true
+freeze_text_model:
+    value: true
+gradient_accumulation_steps:
+    value: 1
+learning_rate:
+    value: 0.0001
+logging_steps:
+    value: 100
+max_length_protein:
+    value: 1024
+max_length_text:
+    value: 512
+num_epochs:
+    value: 10
+num_query_tokens:
+    value: 8
+num_workers:
+    value: 8
+output_dir:
+    value: ./contrastive_outputs
+protein_model_name:
+    value: /nas/shared/kilab/wangyujia/ProtT3/plm_model/esm2-150m
+protein_weight:
+    value: 1
+ptm_weight:
+    value: 1
+qformer_model_name:
+    value: /nas/shared/kilab/wangyujia/ProtT3/plm_model/microsoft
+save_steps:
+    value: 1000
+save_total_limit:
+    value: 3
+seed:
+    value: 42
+temperature:
+    value: 0.07
+text_model_name:
+    value: /oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged
+text_weight:
+    value: 1
+use_wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: protein-llm-contrastive
+warmup_steps:
+    value: 1000
+weight_decay:
+    value: 0.01

BioReason_new/wandb/run-20250811_220309-2qgjwsxa/files/output.log

@@ -0,0 +1,30 @@
+Loading model...
+Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████| 4/4 [00:01<00:00,  2.49it/s]
+Some weights of EsmModel were not initialized from the model checkpoint at /nas/shared/kilab/wangyujia/ProtT3/plm_model/esm2-150m and are newly initialized: ['pooler.dense.bias', 'pooler.dense.weight']
+You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
+Loading datasets...
+Traceback (most recent call last):
+  File "/nas/shared/kilab/wangyujia/BioReason_new/train_contrastive.py", line 549, in <module>
+    trainer = main(args)
+              ^^^^^^^^^^
+  File "/nas/shared/kilab/wangyujia/BioReason_new/train_contrastive.py", line 317, in main
+    train_dataset = load_dataset("json", args.dataset_name, split="train")
+                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/root/miniconda3/envs/bioreason/lib/python3.11/site-packages/datasets/load.py", line 1392, in load_dataset
+    builder_instance = load_dataset_builder(
+                       ^^^^^^^^^^^^^^^^^^^^^
+  File "/root/miniconda3/envs/bioreason/lib/python3.11/site-packages/datasets/load.py", line 1166, in load_dataset_builder
+    builder_instance: DatasetBuilder = builder_cls(
+                                       ^^^^^^^^^^^^
+  File "/root/miniconda3/envs/bioreason/lib/python3.11/site-packages/datasets/builder.py", line 343, in __init__
+    self.config, self.config_id = self._create_builder_config(
+                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/root/miniconda3/envs/bioreason/lib/python3.11/site-packages/datasets/builder.py", line 552, in _create_builder_config
+    builder_config = self.BUILDER_CONFIG_CLASS(**config_kwargs)
+                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "<string>", line 16, in __init__
+  File "/root/miniconda3/envs/bioreason/lib/python3.11/site-packages/datasets/packaged_modules/json/json.py", line 55, in __post_init__
+    super().__post_init__()
+  File "/root/miniconda3/envs/bioreason/lib/python3.11/site-packages/datasets/builder.py", line 126, in __post_init__
+    raise InvalidConfigName(
+datasets.builder.InvalidConfigName: Bad characters from black list '<>:/\|?*' found in '/nas/shared/kilab/wangyujia/ProtT3/data/SwissProtV3/train_set.jsonl'. They could create issues when creating a directory for this config on Windows filesystem.

BioReason_new/wandb/run-20250811_220309-2qgjwsxa/files/requirements.txt

@@ -0,0 +1,233 @@
+nvidia-nccl-cu12==2.26.2
+cbor2==5.6.5
+jupyter_server==2.16.0
+nvidia-curand-cu12==10.3.7.77
+bleach==6.2.0
+py-cpuinfo==9.0.0
+llvmlite==0.44.0
+fsspec==2025.3.0
+uvloop==0.21.0
+rfc3986-validator==0.1.1
+smmap==5.0.2
+pip==25.1
+compressed-tensors==0.10.2
+ipython_pygments_lexers==1.1.1
+fastapi-cli==0.0.8
+filelock==3.18.0
+msgspec==0.19.0
+hjson==3.1.0
+markdown-it-py==3.0.0
+pyzmq==27.0.1
+interegular==0.3.3
+widgetsnbextension==4.0.14
+vllm==0.10.0
+ipykernel==6.30.1
+pydantic==2.11.7
+click==8.2.1
+torchvision==0.22.1
+fastapi-cloud-cli==0.1.5
+httpcore==1.0.9
+nvidia-cuda-nvrtc-cu12==12.6.77
+mdurl==0.1.2
+rich-toolkit==0.15.0
+Pygments==2.19.2
+pure_eval==0.2.3
+types-python-dateutil==2.9.0.20250809
+referencing==0.36.2
+jupyterlab_widgets==3.0.15
+typing-inspection==0.4.1
+stack-data==0.6.3
+jupyter_client==8.6.3
+regex==2025.7.33
+platformdirs==4.3.8
+arrow==1.3.0
+aiosignal==1.4.0
+python-dateutil==2.9.0.post0
+numpy==2.2.6
+jupyter-lsp==2.2.6
+transformers==4.55.0
+mpmath==1.3.0
+six==1.17.0
+python-json-logger==3.3.0
+distro==1.9.0
+partial-json-parser==0.2.1.1.post6
+bitsandbytes==0.46.1
+nvidia-cusparselt-cu12==0.6.3
+pandocfilters==1.5.1
+pexpect==4.9.0
+pydantic-extra-types==2.10.5
+Jinja2==3.1.6
+sentencepiece==0.2.0
+uvicorn==0.35.0
+babel==2.17.0
+trl==0.21.0
+urllib3==2.5.0
+prometheus_client==0.22.1
+watchfiles==1.1.0
+prometheus-fastapi-instrumentator==7.1.0
+jsonschema-specifications==2025.4.1
+diskcache==5.6.3
+webcolors==24.11.1
+peft==0.17.0
+jiter==0.10.0
+triton==3.3.1
+gitdb==4.0.12
+gguf==0.17.1
+safetensors==0.6.2
+cloudpickle==3.1.1
+multiprocess==0.70.16
+aiohttp==3.12.15
+tornado==6.5.2
+nvidia-nvtx-cu12==12.6.77
+nbclient==0.10.2
+nbconvert==7.16.6
+psutil==7.0.0
+llguidance==0.7.30
+ray==2.48.0
+wcwidth==0.2.13
+rignore==0.6.4
+nvidia-cudnn-cu12==9.5.1.17
+soupsieve==2.7
+wandb==0.21.1
+overrides==7.7.0
+opencv-python-headless==4.12.0.88
+pycparser==2.22
+scipy==1.16.1
+terminado==0.18.1
+typer==0.16.0
+parso==0.8.4
+lark==1.2.2
+msgpack==1.1.1
+websockets==15.0.1
+idna==3.10
+fastrlock==0.8.3
+jedi==0.19.2
+accelerate==1.10.0
+jupyter==1.1.1
+beautifulsoup4==4.13.4
+h11==0.16.0
+MarkupSafe==3.0.2
+python-dotenv==1.1.1
+aiohappyeyeballs==2.6.1
+rich==14.1.0
+nbformat==5.10.4
+traitlets==5.14.3
+decorator==5.2.1
+soxr==0.5.0.post1
+propcache==0.3.2
+ninja==1.11.1.4
+cffi==1.17.1
+cupy-cuda12x==13.5.1
+pandas==2.3.1
+deepspeed==0.17.4
+setuptools==78.1.1
+websocket-client==1.8.0
+qwen-vl-utils==0.0.11
+webencodings==0.5.1
+httptools==0.6.4
+jupyterlab==4.4.5
+ptyprocess==0.7.0
+shellingham==1.5.4
+attrs==25.3.0
+fqdn==1.5.1
+huggingface-hub==0.34.4
+tokenizers==0.21.4
+asttokens==3.0.0
+jupyter_server_terminals==0.5.3
+av==15.0.0
+nvidia-cuda-cupti-cu12==12.6.80
+typing_extensions==4.14.1
+hf-xet==1.1.7
+jupyter_core==5.8.1
+starlette==0.47.2
+fastjsonschema==2.21.1
+fastapi==0.116.1
+lightning-utilities==0.15.2
+jupyter-console==6.6.3
+pybase64==1.4.2
+jupyter-events==0.12.0
+requests==2.32.4
+numba==0.61.2
+networkx==3.5
+nvidia-cusparse-cu12==12.5.4.2
+jsonpointer==3.0.0
+pyarrow==21.0.0
+dnspython==2.7.0
+torchaudio==2.7.1
+ipython==9.4.0
+isoduration==20.11.0
+bioreason==0.1.0
+matplotlib-inline==0.1.7
+packaging==25.0
+xxhash==3.5.0
+depyf==0.19.0
+sentry-sdk==2.34.1
+prompt_toolkit==3.0.51
+nvidia-cublas-cu12==12.6.4.1
+rfc3339-validator==0.1.4
+nvidia-cufft-cu12==11.3.0.4
+email_validator==2.2.0
+pycountry==24.6.1
+argon2-cffi==25.1.0
+nvidia-cufile-cu12==1.11.1.6
+frozenlist==1.7.0
+json5==0.12.0
+tinycss2==1.4.0
+defusedxml==0.7.1
+lm-format-enforcer==0.10.12
+Send2Trash==1.8.3
+anyio==4.10.0
+rfc3987-syntax==1.1.0
+pydantic_core==2.33.2
+debugpy==1.8.16
+async-lru==2.0.5
+nvidia-cuda-runtime-cu12==12.6.77
+tiktoken==0.11.0
+comm==0.2.3
+PyYAML==6.0.2
+blake3==1.0.5
+nvidia-cusolver-cu12==11.7.1.2
+torch==2.7.1
+torchmetrics==1.8.1
+yarl==1.20.1
+dill==0.3.8
+wheel==0.45.1
+cachetools==6.1.0
+multidict==6.6.3
+pytz==2025.2
+pillow==11.3.0
+annotated-types==0.7.0
+astor==0.8.1
+nest-asyncio==1.6.0
+httpx==0.28.1
+argon2-cffi-bindings==25.1.0
+notebook_shim==0.2.4
+jsonschema==4.25.0
+python-multipart==0.0.20
+charset-normalizer==3.4.3
+tqdm==4.67.1
+xformers==0.0.31
+tzdata==2025.2
+einops==0.8.1
+mistral_common==1.8.3
+jupyterlab_server==2.27.3
+sympy==1.14.0
+datasets==4.0.0
+GitPython==3.1.45
+mistune==3.1.3
+ipywidgets==8.1.7
+nvidia-ml-py==13.580.65
+uri-template==1.3.0
+notebook==7.4.5
+certifi==2025.8.3
+nvidia-nvjitlink-cu12==12.6.85
+openai==1.90.0
+xgrammar==0.1.21
+executing==2.2.0
+soundfile==0.13.1
+jupyterlab_pygments==0.3.0
+outlines_core==0.2.10
+sniffio==1.3.1
+pytorch-lightning==2.5.2
+rpds-py==0.27.0
+protobuf==6.31.1

BioReason_new/wandb/run-20250811_220309-2qgjwsxa/files/wandb-metadata.json

@@ -0,0 +1,113 @@
+{
+  "os": "Linux-5.10.134-008.18.kangaroo.al8.x86_64-x86_64-with-glibc2.35",
+  "python": "CPython 3.11.0",
+  "startedAt": "2025-08-11T14:03:09.687288Z",
+  "args": [
+    "--text_model_name",
+    "/oss/wangyujia/BIO/construction_finetuning/alpaca/v1-20250609-141541/checkpoint-50-merged",
+    "--protein_model_name",
+    "/nas/shared/kilab/wangyujia/ProtT3/plm_model/esm2-150m",
+    "--qformer_model_name",
+    "/nas/shared/kilab/wangyujia/ProtT3/plm_model/microsoft",
+    "--num_query_tokens",
+    "8",
+    "--dataset_name",
+    "/nas/shared/kilab/wangyujia/ProtT3/data/SwissProtV3/train_set.jsonl",
+    "--output_dir",
+    "./contrastive_outputs",
+    "--num_epochs",
+    "10",
+    "--batch_size",
+    "32",
+    "--learning_rate",
+    "1e-4",
+    "--temperature",
+    "0.07",
+    "--freeze_protein_model",
+    "--freeze_text_model",
+    "--enable_ptm",
+    "--max_length_protein",
+    "1024",
+    "--max_length_text",
+    "512",
+    "--num_workers",
+    "8",
+    "--eval_dataset",
+    "--use_wandb",
+    "--wandb_project",
+    "protein-llm-contrastive",
+    "--logging_steps",
+    "100",
+    "--eval_steps",
+    "500",
+    "--save_steps",
+    "1000"
+  ],
+  "program": "/nas/shared/kilab/wangyujia/BioReason_new/train_contrastive.py",
+  "codePath": "wangyujia/BioReason_new/train_contrastive.py",
+  "codePathLocal": "train_contrastive.py",
+  "git": {
+    "remote": "https://github.com/PorUna-byte/PAR.git",
+    "commit": "b8caf406aa1699c788f0ca6e44a1769452c317db"
+  },
+  "root": "/nas/shared/kilab/wangyujia/BioReason_new",
+  "host": "dsw-265304-f8bc5ff76-4mdt5",
+  "executable": "/root/miniconda3/envs/bioreason/bin/python",
+  "cpu_count": 64,
+  "cpu_count_logical": 64,
+  "gpu": "NVIDIA A800-SXM4-80GB",
+  "gpu_count": 8,
+  "disk": {
+    "/": {
+      "total": "1623302262784",
+      "used": "28193923072"
+    }
+  },
+  "memory": {
+    "total": "549755813888"
+  },
+  "gpu_nvidia": [
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "architecture": "Ampere",
+      "uuid": "GPU-71607f78-ad31-1ea4-19c1-908e3e31aaf1"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "architecture": "Ampere",
+      "uuid": "GPU-92b7dbbd-7ef5-3c5f-ce1c-1d179d7fa587"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "architecture": "Ampere",
+      "uuid": "GPU-bbc35439-ad79-578b-381b-aba6f0cc0168"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "architecture": "Ampere",
+      "uuid": "GPU-e492e147-ca2e-76f2-85da-4e08e4deeb14"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "architecture": "Ampere",
+      "uuid": "GPU-8c4f8e67-4b52-5107-3095-0f007e6378ac"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "architecture": "Ampere",
+      "uuid": "GPU-7063f0b9-4ca2-6a72-522d-1262899ac5ad"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "architecture": "Ampere",
+      "uuid": "GPU-3b6e9a37-bcf3-387c-7874-4f8de4abd115"
+    },
+    {
+      "name": "NVIDIA A800-SXM4-80GB",
+      "architecture": "Ampere",
+      "uuid": "GPU-92456839-e814-7be9-6817-f3e8da8aa80c"
+    }
+  ],
+  "cudaVersion": "12.1",
+  "writerId": "r185oiuz6xjarzg7yyfap3b9flv6ll88"
+}