text_tokenizer.py

import pathlib
from typing import Optional, List, Union
import torch
from transformers import BertTokenizer

class BioBERTTokenizerWrapper:
    def __init__(
        self,
        model_path: str = None,
        max_length: int = 64,
    ):
        # 1. 确定路径
        if model_path is None:
            base = pathlib.Path(__file__).parent.parent.parent.resolve()
            model_path = str(base / 'data' / 'biobert')
        
        self.model_path = str(model_path)
        self.max_length = max_length

        # 2.加载 BioBERT Tokenizer (使用 transformers)
        print(f"📥 正在从 {self.model_path} 加载 BioBERT Tokenizer...")
        self.tokenizer = BertTokenizer.from_pretrained(self.model_path)
        
        # 3. 获取 pad_token_id（BioBERT 使用 0 作为 pad_token_id）
        self.pad_token_id = self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else 0
        
        # 4. 获取词表大小
        self.vocab_size = self.tokenizer.vocab_size
        print(f"✅ BioBERT Tokenizer 就绪。词表大小: {self.vocab_size}, 最大长度: {self.max_length}")

    def __call__(self, text: Union[str, List[str]], return_tensors: str = 'pt', padding: bool = True, truncation: bool = True, max_length: Optional[int] = None) -> dict:
        """
        输入: 文本 或 文本列表
        输出: 包含 input_ids 和 attention_mask 的字典
        """
        if max_length is None:
            max_length = self.max_length
            
        if not text:
            # 如果文本为空，返回全 Pad 的序列
            batch_size = 1
            input_ids = torch.full((batch_size, max_length), self.pad_token_id, dtype=torch.long)
            attention_mask = torch.zeros((batch_size, max_length), dtype=torch.long)
            return {'input_ids': input_ids, 'attention_mask': attention_mask}
            
        if isinstance(text, str):  # 单个字符串，转换为列表
            text = [text]

        # 使用 transformers 的 tokenizer 进行编码
        encoded = self.tokenizer(
            text,
            padding='max_length' if padding else False,
            truncation=truncation,
            max_length=max_length,
            return_tensors=return_tensors
        )
        
        return encoded
    
    def tokenize(self, text: Union[str, List[str]]) -> torch.Tensor:
        """
        输入: 文本 或 文本列表
        输出: Token IDs (LongTensor) [Batch, Max_Len]
        """
        result = self(text, return_tensors='pt', padding=True, truncation=True)
        return result['input_ids']

    def detokenize(self, token_ids: torch.Tensor) -> List[str]:
        """用于调试：查看 ID 对应的文本"""
        if isinstance(token_ids, torch.Tensor):
            token_ids = token_ids.cpu().numpy().tolist()
        elif isinstance(token_ids, list):
            pass
        else:
            token_ids = token_ids.tolist()
        
        # 使用 transformers 的 decode 方法
        decoded_texts = self.tokenizer.batch_decode(token_ids, skip_special_tokens=True)
        return decoded_texts