biobert_encoder.py

import pathlib
from typing import Optional, List, Tuple
import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer, BertConfig


def _resolve_biobert_path(biobert_path: Optional[str]) -> str:
    """Resolve a BioBERT checkpoint path against the repository layout."""
    current_file = pathlib.Path(__file__).resolve()
    repo_root = current_file.parent.parent.parent

    candidate_paths = []
    if biobert_path:
        raw_path = pathlib.Path(biobert_path)
        candidate_paths.append(raw_path)
        if not raw_path.is_absolute():
            candidate_paths.append((repo_root / raw_path).resolve())

    candidate_paths.extend(
        [
            repo_root / "checkpoints" / "biobert",
            repo_root / "data" / "biobert",
        ]
    )

    for candidate in candidate_paths:
        candidate = pathlib.Path(candidate)
        if candidate.exists():
            return str(candidate.resolve())

    return str(pathlib.Path(biobert_path).resolve()) if biobert_path else str((repo_root / "checkpoints" / "biobert").resolve())

class BioBERTFunctionEncoder(nn.Module):
    def __init__(
        self,
        biobert_path: str = None,
        max_length: int = 128,
        model_dim: int = 768,
        device: Optional[str] = None,
    ):
        super().__init__()
        
        self.biobert_path = _resolve_biobert_path(biobert_path)
        self.max_length = max_length
        self.target_dim = model_dim # 目标 768

        # 1. 加载模型和分词器
        if device is None:
            device_obj = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            device_obj = torch.device(device)
        self._model_device = device_obj  # 内部使用 _model_device 跟踪实际设备
        
        # 使用 transformers 加载 BioBERT 模型和 tokenizer
        print(f"📥 正在从 {self.biobert_path} 加载 BioBERT 模型...")
        self.tokenizer = BertTokenizer.from_pretrained(self.biobert_path)
        config = BertConfig.from_pretrained(self.biobert_path)
        self.model = BertModel(config)
        self.model.to(self._model_device)
        self.model.eval()  # 设为评估模式，不参与训练
    
        # 2.检查输出维度
        # BioBERT 的隐藏层维度是 768
        self.raw_output_dim = self.model.config.hidden_size
            
        print(f"✅ BioBERT 功能文本编码器: {self.biobert_path} (Raw dim={self.raw_output_dim}, Final dim={self.target_dim}, 序列输出)")

    @property
    def device(self):
        """返回模型当前所在的设备"""
        return self._model_device
    
    @device.setter
    def device(self, value):
        """设置模型设备并移动模型"""
        if isinstance(value, str):
            value = torch.device(value)
        self._model_device = value
        # 只有在 model 已经创建后才移动
        if hasattr(self, 'model') and self.model is not None:
            self.model.to(value)

    def encode_batch(self, texts: List[str], batch_size: int = 8, target_device: Optional[torch.device] = None) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        批量编码文本，返回序列隐藏状态和对应的 attention mask。
        Output Shape: (Hidden_States [B, L_text, D], Attention_Mask [B, L_text])
        
        Args:
            texts: 文本列表
            batch_size: 批处理大小
            target_device: 目标设备，如果为 None 则使用 self.device（GPU 加速）
        """
        # 确定实际使用的设备
        if target_device is not None:
            actual_device = target_device
        else:
            actual_device = self._model_device
        
        if not texts:
            return torch.zeros(0, self.max_length, self.target_dim, device=actual_device), torch.zeros(0, self.max_length, dtype=torch.long, device=actual_device)
            
        batched_hiddens = []
        batched_masks = []
        
        for start in range(0, len(texts), batch_size):
            batch_texts = [t or '' for t in texts[start:start + batch_size]]
            
            # 使用 tokenizer 进行编码
            encoded = self.tokenizer(
                batch_texts,
                padding='max_length',
                truncation=True,
                max_length=self.max_length,
                return_tensors='pt'
            )
            
            # 移动到模型设备
            input_ids = encoded['input_ids'].to(self._model_device)
            attention_mask = encoded['attention_mask'].to(self._model_device)
            
            # 使用模型进行编码
            with torch.no_grad():
                outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
                # outputs.last_hidden_state: [B, L, 768]
                hidden_states = outputs.last_hidden_state
            
            # 🔧 如果目标设备与模型设备不同，则移动结果
            if target_device is not None and target_device != self._model_device:
                hidden_states = hidden_states.to(target_device)
                attention_mask = attention_mask.to(target_device)
            
            batched_hiddens.append(hidden_states)
            batched_masks.append(attention_mask)
        
        # 堆叠所有 batch
        if batched_hiddens:
            return torch.cat(batched_hiddens, dim=0), torch.cat(batched_masks, dim=0)
        else:
            return torch.zeros(0, self.max_length, self.target_dim, device=actual_device), torch.zeros(0, self.max_length, dtype=torch.long, device=actual_device)

    #  新增 forward 方法，作为 BioBERTFunctionEncoder 的主要接口
    def forward(self, texts: List[str]) -> Tuple[torch.Tensor, torch.Tensor]:
        """训练和验证时调用的接口"""
        return self.encode_batch(texts)