easy_tpp/preprocess/robert_tokenizer.py

"""
评论罗伯特事件分词器 - 支持语义特征、偏差特征等

扩展EventTokenizer以支持多模态特征的padding和批处理
"""

from typing import Optional, Union, Dict, Any
import numpy as np
import torch
from easy_tpp.preprocess.event_tokenizer import EventTokenizer, BatchEncoding
from easy_tpp.utils import PaddingStrategy


class RobertEventTokenizer(EventTokenizer):
    """
    支持语义特征、偏差特征等的事件分词器
    
    扩展EventTokenizer以支持：
    - semantic_vectors: 语义向量padding
    - deviation_features: 偏差特征padding
    - is_spontaneous: 自发/被@标记padding
    """
    
    def __init__(self, config, use_semantic=False, use_deviation=False, semantic_dim=768):
        """
        初始化分词器
        
        Args:
            config: 配置对象
            use_semantic: 是否使用语义特征
            use_deviation: 是否使用偏差特征
            semantic_dim: 语义向量维度
        """
        super(RobertEventTokenizer, self).__init__(config)
        
        self.use_semantic = use_semantic
        self.use_deviation = use_deviation
        self.semantic_dim = semantic_dim
        
        # 添加自定义特征到model_input_names
        # 标准顺序：time_seqs, time_delta_seqs, type_seqs, batch_non_pad_mask, attention_mask
        # 自定义特征添加在后面
        if self.use_semantic:
            self.model_input_names.append('semantic_vectors')
        if self.use_deviation:
            self.model_input_names.append('deviation_features')
        
        # is_spontaneous总是添加（如果使用）
        self.model_input_names.append('is_spontaneous')
    
    def _pad(
        self,
        encoded_inputs: Union[Dict[str, Any], BatchEncoding],
        max_length: Optional[int] = None,
        padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
        return_attention_mask: Optional[bool] = None,
    ) -> dict:
        """
        填充编码输入（包括自定义特征）
        
        Args:
            encoded_inputs: 编码后的输入
            max_length: 最大长度
            padding_strategy: 填充策略
            return_attention_mask: 是否返回注意力掩码
        
        Returns:
            dict: 填充后的批次数据
        """
        # 先处理标准字段（调用父类方法）
        # 但我们需要重写以添加自定义特征的处理
        required_input = encoded_inputs[self.model_input_names[0]]
        
        if padding_strategy == PaddingStrategy.LONGEST:
            max_length = max(len(seq) for seq in required_input)
            padding_strategy = PaddingStrategy.MAX_LENGTH
        
        # 获取序列长度
        seq_lens = np.array([len(seq) for seq in required_input])
        is_all_seq_equal_max_length = np.all(seq_lens == max_length)
        needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and ~is_all_seq_equal_max_length
        
        batch_output = dict()
        
        # 处理标准字段（time_seqs, time_delta_seqs, type_seqs）
        if needs_to_be_padded:
            batch_output[self.model_input_names[0]] = self.make_pad_sequence(
                encoded_inputs[self.model_input_names[0]],
                self.pad_token_id,
                padding_side=self.padding_side,
                max_len=max_length
            )
            batch_output[self.model_input_names[1]] = self.make_pad_sequence(
                encoded_inputs[self.model_input_names[1]],
                self.pad_token_id,
                padding_side=self.padding_side,
                max_len=max_length
            )
            batch_output[self.model_input_names[2]] = self.make_pad_sequence(
                encoded_inputs[self.model_input_names[2]],
                self.pad_token_id,
                padding_side=self.padding_side,
                max_len=max_length,
                dtype=np.int64
            )
        else:
            batch_output[self.model_input_names[0]] = np.array(
                encoded_inputs[self.model_input_names[0]], dtype=np.float32
            )
            batch_output[self.model_input_names[1]] = np.array(
                encoded_inputs[self.model_input_names[1]], dtype=np.float32
            )
            batch_output[self.model_input_names[2]] = np.array(
                encoded_inputs[self.model_input_names[2]], dtype=np.int64
            )
        
        # non_pad_mask
        seq_pad_mask = np.full_like(batch_output[self.model_input_names[2]], fill_value=True, dtype=bool)
        for i, seq_len in enumerate(seq_lens):
            seq_pad_mask[i, seq_len:] = False
        batch_output[self.model_input_names[3]] = seq_pad_mask
        
        # attention_mask
        if return_attention_mask is None:
            return_attention_mask = "attention_mask" in self.model_input_names
        
        if return_attention_mask:
            batch_output[self.model_input_names[4]] = self.make_attn_mask_for_pad_sequence(
                batch_output[self.model_input_names[2]],
                self.pad_token_id
            )
        else:
            batch_output[self.model_input_names[4]] = []
        
        # 处理自定义特征
        # 处理语义向量
        if self.use_semantic and 'semantic_vectors' in encoded_inputs:
            semantic_vectors = encoded_inputs['semantic_vectors']
            if needs_to_be_padded:
                batch_output['semantic_vectors'] = self.make_pad_sequence_for_features(
                    semantic_vectors,
                    pad_value=0.0,
                    max_len=max_length,
                    feature_dim=self.semantic_dim
                )
            else:
                batch_output['semantic_vectors'] = np.array(semantic_vectors, dtype=np.float32)
        elif self.use_semantic:
            # 如果没有提供但需要，创建零向量
            batch_size = len(required_input)
            if needs_to_be_padded:
                batch_output['semantic_vectors'] = np.zeros(
                    (batch_size, max_length, self.semantic_dim), dtype=np.float32
                )
            else:
                # 使用最大长度
                max_seq_len = int(seq_lens.max())
                batch_output['semantic_vectors'] = np.zeros(
                    (batch_size, max_seq_len, self.semantic_dim), dtype=np.float32
                )
        
        # 处理偏差特征
        if self.use_deviation and 'deviation_features' in encoded_inputs:
            deviation_features = encoded_inputs['deviation_features']
            if needs_to_be_padded:
                batch_output['deviation_features'] = self.make_pad_sequence_for_features(
                    deviation_features,
                    pad_value=0.0,
                    max_len=max_length,
                    feature_dim=3
                )
            else:
                batch_output['deviation_features'] = np.array(deviation_features, dtype=np.float32)
        elif self.use_deviation:
            # 如果没有提供但需要，创建零向量
            batch_size = len(required_input)
            if needs_to_be_padded:
                batch_output['deviation_features'] = np.zeros(
                    (batch_size, max_length, 3), dtype=np.float32
                )
            else:
                max_seq_len = int(seq_lens.max())
                batch_output['deviation_features'] = np.zeros(
                    (batch_size, max_seq_len, 3), dtype=np.float32
                )
        
        # 处理is_spontaneous
        if 'is_spontaneous' in encoded_inputs:
            is_spontaneous = encoded_inputs['is_spontaneous']
            if needs_to_be_padded:
                batch_output['is_spontaneous'] = self.make_pad_sequence_for_features(
                    is_spontaneous,
                    pad_value=-1.0,  # -1表示不适用
                    max_len=max_length,
                    feature_dim=1  # 标量
                )
            else:
                batch_output['is_spontaneous'] = np.array(is_spontaneous, dtype=np.float32)
        else:
            # 如果没有提供，创建-1向量（不适用）
            batch_size = len(required_input)
            if needs_to_be_padded:
                batch_output['is_spontaneous'] = np.full(
                    (batch_size, max_length), -1.0, dtype=np.float32
                )
            else:
                max_seq_len = int(seq_lens.max())
                batch_output['is_spontaneous'] = np.full(
                    (batch_size, max_seq_len), -1.0, dtype=np.float32
                )
        
        return batch_output
    
    def make_pad_sequence_for_features(self, seqs, pad_value, max_len, feature_dim=None, dtype=np.float32):
        """
        为特征序列创建padding（辅助方法）
        
        Args:
            seqs: 序列列表
            pad_value: padding值
            max_len: 最大长度
            feature_dim: 特征维度（如果是多维特征）
            dtype: 数据类型
        
        Returns:
            np.ndarray: 填充后的数组
        """
        padded_seqs = []
        for seq in seqs:
            seq_len = len(seq)
            if seq_len < max_len:
                pad_len = max_len - seq_len
                if isinstance(seq, np.ndarray):
                    if seq.ndim == 1:
                        # 一维数组
                        pad = np.full(pad_len, pad_value, dtype=dtype)
                        padded_seq = np.concatenate([seq, pad], axis=0)
                    else:
                        # 多维数组
                        pad_shape = (pad_len,) + seq.shape[1:]
                        pad = np.full(pad_shape, pad_value, dtype=dtype)
                        padded_seq = np.concatenate([seq, pad], axis=0)
                else:
                    # 列表
                    if isinstance(seq[0], (list, np.ndarray, tuple)):
                        # 嵌套列表（多维）
                        if feature_dim:
                            pad = [[pad_value] * feature_dim] * pad_len
                        else:
                            pad = [[pad_value] * len(seq[0])] * pad_len
                        padded_seq = seq + pad
                    else:
                        # 一维列表
                        pad = [pad_value] * pad_len
                        padded_seq = seq + pad
                padded_seqs.append(padded_seq)
            else:
                padded_seqs.append(seq[:max_len])
        
        return np.array(padded_seqs, dtype=dtype)