Upload 2 files

DataProcessor_pdeeppp.py

@@ -0,0 +1,87 @@
+from transformers.processing_utils import ProcessorMixin
+from transformers.tokenization_utils_base import BatchEncoding
+
+
+class PDeepPPProcessor(ProcessorMixin):
+    def __init__(self, pad_char="X", target_length=33):
+        self.pad_char = pad_char
+        self.target_length = target_length
+    
+    def pad_sequence(self, seq):
+        """确保序列长度为 target_length，不足的部分用 pad_char 在两侧均匀填充"""
+        if len(seq) < self.target_length:
+            total_padding = self.target_length - len(seq)
+            left_padding = total_padding // 2
+            right_padding = total_padding - left_padding
+            seq = self.pad_char * left_padding + seq + self.pad_char * right_padding
+        return seq[:self.target_length]
+    
+    def extract_ptm_sequences(self, sequences):
+        """处理 PTM 数据，确保目标氨基酸（S、T、Y）位于序列中心"""
+        ptm_data = []
+        for seq in sequences:
+            for i in range(len(seq)):
+                if seq[i] in {'S', 'T', 'Y'}:  # 仅提取 S、T、Y 作为中心的片段
+                    start = max(0, i - self.target_length // 2)
+                    end = min(len(seq), start + self.target_length)
+                    padded_seq = self.pad_sequence(seq[start:end])
+                    ptm_data.append(padded_seq)
+        return ptm_data
+    
+    def extract_bps_sequences(self, sequences, overlapping=True, step_size=5):
+        """处理生物活性数据（BPS），关注整个序列，可重叠"""
+        bioactive_data = []
+        for seq in sequences:
+            if len(seq) < self.target_length:
+                # 如果序列长度不足，直接填充到 target_length
+                padded_seq = self.pad_sequence(seq)
+                bioactive_data.append(padded_seq)
+            else:
+                # 如果序列长度足够，按照滑动窗口提取片段
+                for i in range(0, len(seq) - self.target_length + 1, 
+                            step_size if overlapping else self.target_length):
+                    bioactive_data.append(self.pad_sequence(seq[i:i + self.target_length]))
+        return bioactive_data
+    
+    def __call__(
+        self, 
+        sequences, 
+        mode,  # 去除默认值，强制外部传入
+        overlapping=True,
+        step_size=5,
+        **kwargs
+    ):
+        """
+        预处理蛋白质序列，仅处理数据到指定长度。
+
+        Args:
+            sequences: 序列列表或单个序列字符串。
+            mode: 选择处理模式，必须从外部传入，"PTM" 或 "BPS"。
+            overlapping: BPS 模式下是否使用重叠窗口。
+            step_size: BPS 模式下的步长。
+        """
+        # 确保 sequences 是列表
+        if isinstance(sequences, str):
+            sequences = [sequences]
+            
+        # 根据模式提取序列
+        if mode == "PTM":
+            processed_sequences = self.extract_ptm_sequences(sequences)
+        elif mode == "BPS":
+            processed_sequences = self.extract_bps_sequences(
+                sequences, 
+                overlapping=overlapping,
+                step_size=step_size
+            )
+        else:
+            raise ValueError("Invalid mode. Please choose 'PTM' or 'BPS'.")
+        
+        if len(processed_sequences) == 0:
+            raise ValueError("No sequences processed. Check input data and processing logic.")
+
+        # 创建返回字典，仅包含预处理后的序列
+        model_inputs = {
+            "raw_sequences": processed_sequences,  # 预处理后的序列
+        }
+
+        return BatchEncoding(data=model_inputs)  # 返回处理后的数据

Pretraining_pdeeppp.py

@@ -0,0 +1,119 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class PretrainingPDeepPP:
+    def __init__(self, embedding_dim=1280, target_length=33, esm_ratio=None, device=None):
+        """
+        初始化 PretrainingPDeepPP 类。
+
+        Args:
+            embedding_dim: 嵌入维度大小。
+            target_length: 目标序列长度。
+            esm_ratio: ESM 表征与嵌入表示的权重比例（由外部赋值）。
+            device: 设备信息。
+        """
+        self.embedding_dim = embedding_dim
+        self.target_length = target_length
+        self.esm_ratio = esm_ratio  # 仅存储 esm_ratio，不赋默认值
+        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def extract_esm_representations(self, sequences, esm_model, batch_converter, batch_size=32):
+        """
+        提取 ESM 表征，并直接返回形状为 (batch_size, target_length, embedding_dim) 的结果。
+        """
+        sequence_representations = []
+        print("Sequences to process:", sequences)
+        print("Batch size:", batch_size)
+        
+        # 为每个序列添加一个“伪标签”以满足 batch_converter 要求
+        labeled_sequences = [(None, seq) for seq in sequences]
+
+        for i in range(0, len(labeled_sequences), batch_size):
+            batch = labeled_sequences[i:i + batch_size]
+            if len(batch) == 0:
+                continue
+            # 调用 batch_converter 将序列转换为 batch_tokens
+            _, batch_strs, batch_tokens = batch_converter(batch)
+            batch_tokens = batch_tokens.to(self.device)
+            
+            # 使用 ESM 模型提取表示
+            with torch.no_grad():
+                results = esm_model(batch_tokens, repr_layers=[33], return_contacts=False)
+            
+            # 提取每个序列的表示
+            for token_repr in results["representations"][33]:  # 获取第 33 层的表示
+                sequence_representations.append(token_repr[:self.target_length])
+
+        if len(sequence_representations) == 0:
+            raise ValueError("No ESM representations were generated. Check your input sequences and batch processing logic.")
+
+        # 将所有序列的表示堆叠起来，形状为 (batch_size, 33, 1280)
+        return torch.stack(sequence_representations)
+
+    def pad_sequences(self, sequences, max_len=None, pad_value=0):
+        if max_len is None:
+            max_len = max(len(seq) for seq in sequences)
+        padded_sequences = torch.zeros((len(sequences), max_len), dtype=torch.long)
+        for i, seq in enumerate(sequences):
+            padded_sequences[i, :len(seq)] = torch.tensor(seq)
+        return padded_sequences
+
+    def seq_to_indices(self, seq, vocab_dict):
+        return [vocab_dict.get(char, 0) for char in seq]
+
+    def create_embeddings(self, sequences, vocab, esm_model, esm_alphabet, batch_size=16):
+        """
+        创建嵌入向量，使用类的 esm_ratio 属性动态控制权重分配。
+
+        Args:
+            sequences: 输入序列列表。
+            vocab: 字符词汇表。
+            esm_model: 预训练的 ESM 模型。
+            esm_alphabet: ESM 模型的字母表。
+            batch_size: 批量大小。
+
+        Returns:
+            结合 ESM 表征与嵌入表示的嵌入结果。
+        """
+        if self.esm_ratio is None:
+            raise ValueError("esm_ratio is not set. Please assign a value before creating embeddings.")
+
+        # 构建词汇表字典
+        vocab_dict = {char: i for i, char in enumerate(vocab)}
+
+        # 将序列转为索引
+        indices = [self.seq_to_indices(seq, vocab_dict) for seq in sequences]
+        indices_padded = self.pad_sequences(indices, max_len=self.target_length)
+
+        # 定义嵌入模型
+        class EmbeddingPretrainedModel(nn.Module):
+            def __init__(self, vocab_size, embedding_dim, max_len):
+                super(EmbeddingPretrainedModel, self).__init__()
+                self.embedding = nn.Embedding(vocab_size, embedding_dim)
+                self.fc = nn.Linear(embedding_dim, embedding_dim)
+            
+            def forward(self, x):
+                x = self.embedding(x)
+                x = self.fc(x)
+                return x
+
+        embedding_model = EmbeddingPretrainedModel(len(vocab), self.embedding_dim, self.target_length).to(self.device)
+
+        # 提取 ESM 表示
+        esm_representations = self.extract_esm_representations(
+            sequences,
+            esm_model,
+            esm_alphabet.get_batch_converter(),
+            batch_size=batch_size
+        )
+
+        # 获取嵌入表示
+        with torch.no_grad():
+            embedding_output = embedding_model(indices_padded.to(self.device))
+
+        # 合并 ESM 和嵌入表示，动态使用 esm_ratio
+        combined_representations = self.esm_ratio * esm_representations + (1 - self.esm_ratio) * embedding_output
+
+        return combined_representations