dataset/dataset.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Time Series Question Answering Dataset.
Handles loading and preprocessing of time series data and question-answer pairs.
"""
import sys
from transformers import PretrainedConfig, AutoTokenizer
from transformers import AutoProcessor
import torch
import os
import json
from torch.utils.data import Dataset
from typing import List, Dict, Any
import pandas as pd
import numpy as np
import h5py
import re
import random
from models.TimeLanguageModel import TLMConfig
from utils.log_util import adaptive_print

class PretrainDataset(Dataset):

    def __init__(self, ts_path):
        super().__init__()
        self.ts_path = ts_path
        self.load_data()

    def load_data(self):
        with h5py.File(self.ts_path, 'r') as f:
            data = f['seq_data'][:]
        self.datas = data
        
    def __len__(self):
        return len(self.datas)

    def __getitem__(self, index):
        return {'ts_values': torch.tensor(self.datas[index], dtype=torch.float)}

def find_assistant_tokens(tokenizer, target):
    """Find assistant token positions in the target sequence.
    
    Args:
        tokenizer: Tokenizer instance
        target: Target token sequence
        
    Returns:
        List of tuples containing start and end positions of assistant tokens
    """
    result = []
    start_index = 0
    end_index = 0
    while start_index <= len(target) - 1:
        if target[start_index] != tokenizer('assistant')['input_ids'][0]:
            start_index += 1
            end_index += 1
        else:
            end_index += 1
            if target[end_index] == tokenizer('<|im_end|>')['input_ids'][0]:
                result.append((start_index + 1, end_index + 1))
                start_index = end_index + 1
    return result


class TsQaDataset(Dataset):
    """Time Series Question Answering Dataset with token ID range validation."""
    
    def __init__(self, ts_path, data_path, tokenizer, processor, config, pretrain=False, sft=False, shuffle=False):
        """Initialize the dataset.
        
        Args:
            ts_path: Path to time series data file
            data_path: Path to question-answer data file
            tokenizer: Tokenizer instance
            processor: Processor instance
            config: Configuration object
            pretrain: Whether in pretraining mode
            sft: Whether in supervised fine-tuning mode
            shuffle: Whether to shuffle the data
        """
        super().__init__()
        self.ts_path = ts_path
        self.data_path = data_path
        self.tokenizer = tokenizer
        self.processor = processor
        self.config = config
        self.pretrain = pretrain
        self.sft = sft
        self.shuffle = shuffle
        self.h5_file = None
        
        # Key fix: Ensure vocab_size is correct
        self.vocab_size = len(self.tokenizer)
        adaptive_print(f"📊 Vocab size: {self.vocab_size}")
        
        # Ensure tokenizer settings are consistent
        self.tokenizer.padding_side = 'left'
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
            
        # Validate special tokens
        self._validate_special_tokens()
        self._build_index()

    def _validate_special_tokens(self):
        """Validate that all special token IDs are within valid range."""
        special_tokens = {
            'pad_token_id': self.tokenizer.pad_token_id,
            'eos_token_id': self.tokenizer.eos_token_id,
            'bos_token_id': getattr(self.tokenizer, 'bos_token_id', None),
            'unk_token_id': getattr(self.tokenizer, 'unk_token_id', None),
        }
        
        adaptive_print("🔍 Validating special tokens:")
        for name, token_id in special_tokens.items():
            if token_id is not None:
                if token_id >= self.vocab_size or token_id < 0:
                    adaptive_print(f"❌ {name} = {token_id} out of range [0, {self.vocab_size})")
                    # Fix invalid special tokens
                    if name == 'pad_token_id':
                        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
                        adaptive_print(f"🔧 Fixed: pad_token_id -> {self.tokenizer.pad_token_id}")
                else:
                    adaptive_print(f"✅ {name} = {token_id}")

    def _validate_token_ids(self, token_ids, context=""):
        """Validate token IDs for validity.
        
        Args:
            token_ids: List of token IDs to validate
            context: Context string for error messages
            
        Returns:
            List of validated token IDs
        """
        if not isinstance(token_ids, list):
            return token_ids
            
        valid_ids = []
        for i, token_id in enumerate(token_ids):
            if token_id < 0 or token_id >= self.vocab_size:
                adaptive_print(f"⚠️ {context} position {i}: invalid token_id {token_id}, replacing with unk_token")
                # Replace with unk_token, if not available use eos_token
                replacement = getattr(self.tokenizer, 'unk_token_id', self.tokenizer.eos_token_id)
                valid_ids.append(replacement)
            else:
                valid_ids.append(token_id)
        return valid_ids

    def _build_index(self):
        """Build dataset index by loading and processing data files."""
        self.datas = []
        with open(self.data_path, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f):
                item = json.loads(line)
                for i in range(0, len(item['conversations']), 2):
                    if item['conversations'][i]['stage'] in ['1', '2', '3', '4']:
                        self.datas.append({
                            'id': item['id'],
                            'stage': int(item['conversations'][i]['stage']),
                            'form': item['conversations'][i]['attribute'],
                            'question': item['conversations'][i]['value'],
                            'answer': item['conversations'][i + 1]['value'],
                            'line_num': line_num
                        })
        
        if self.shuffle:
            adaptive_print(f"🎲 Shuffling dataset: {self.data_path}")
            random.shuffle(self.datas)

    def _get_h5_file(self):
        """Get HDF5 file handle for time series data."""
        if self.h5_file is None and os.path.exists(self.ts_path):
            self.h5_file = h5py.File(self.ts_path, 'r')
        return self.h5_file

    def __len__(self):
        """Return dataset length."""
        return len(self.datas)

    def add_adaptive_prompt(self, sample):
        sample = sample.copy()
        
        if sample['stage'] == 1:
            sample['question'] += " Please analyze the change in this signal and explain its physical implication, such as component load, airflow, or temperature stability."
        elif sample['stage'] == 2:
            sample['question'] += " Carefully analyze the signal pattern (e.g., stability, oscillation, drops) to determine the correct fault status or root cause. Select the most likely option based on observed signal behavior."
        elif sample['stage'] == 3:
            sample['question'] += " Review the trends across 10 cycles and evaluate the degradation pattern. Select the option that best reflects the long-term health status or risk level indicated by the signal."
        elif sample['stage'] == 4:
            sample['question'] += " Based on the 10-cycle degradation pattern, propose concrete maintenance actions (e.g., replace, inspect) to ensure safe and efficient operation."
        return sample

    def _create_chat_input(self, question):
        """Unified chat input creation method."""
        messages = [
            {"role": "system", "content": 'You are a helpful assistant.'},
            {"role": "user", "content": question}
        ]
        
        try:
            # Use a safer tokenization method
            chat_text = self.tokenizer.apply_chat_template(
                messages, 
                tokenize=False, 
                add_generation_prompt=True
            )
            # Replace time series placeholder
            chat_text = chat_text.replace('<ts>', '<|image_pad|>' * self.config.ts_pad_num)
            return chat_text
        except Exception as e:
            adaptive_print(f"❌ Chat template error: {e}")
            # Fallback to a simple format
            return f"You are a helpful assistant.\nuser\n{question}\nassistant\n"

    def _safe_tokenize(self, text, add_special_tokens=True):
        """Safe tokenization, ensure results are within valid range."""
        try:
            # Add more tokenization parameters
            result = self.tokenizer(
                text, 
                add_special_tokens=add_special_tokens,
                padding=False,
                truncation=False,
                return_tensors=None
            )
            token_ids = result['input_ids']
            
            # Validate token_ids
            token_ids = self._validate_token_ids(token_ids, f"tokenize: {text[:50]}...")
            return token_ids
            
        except Exception as e:
            adaptive_print(f"❌ Tokenization error for text: {text[:100]}...")
            adaptive_print(f"Error: {e}")
            # Return a safe default value
            return [self.tokenizer.eos_token_id]

    def __getitem__(self, idx):
        try:
            sample = self.datas[idx]
            # sample = self.add_adaptive_prompt(sample)

            # Load time series data
            h5f = self._get_h5_file()
            if isinstance(sample['id'], str):
                ts = h5f['seq_data'][int(sample['id']) - 1]
            elif isinstance(sample['id'], list):
                ts_list = [h5f['seq_data'][int(i) - 1][:len(h5f['seq_data'][int(i) - 1]) // 10] for i in sample['id']]
                ts = np.concatenate(ts_list, axis=0)

            # =========================== Mode 1: Pretraining ===========================
            if getattr(self, 'pretrain', False):
                return {
                    'ts_values': torch.tensor(ts, dtype=torch.float)
                }

            # =========================== Mode 2: SFT Training ===========================
            elif getattr(self, 'sft', False):
                # Create query_ids: only the original question text, no other information
                original_question = sample['question']
                query_ids = self._safe_tokenize(original_question, add_special_tokens=False)
                
                # Create input_ids: full input including time series placeholder
                q_text = self._create_chat_input(sample['question'])  # This includes <|image_pad|>
                q_input_ids = self._safe_tokenize(q_text, add_special_tokens=False)
                
                # Ensure answer format is consistent and safe
                a_text = sample['answer']
                if not a_text.endswith(self.tokenizer.eos_token):
                    a_text += self.tokenizer.eos_token
                a_input_ids = self._safe_tokenize(a_text, add_special_tokens=False)

                # Construct training data
                input_ids = q_input_ids + a_input_ids
                labels = [self.tokenizer.pad_token_id] * len(q_input_ids) + a_input_ids

                # Final validation
                query_ids = self._validate_token_ids(query_ids, f"query_ids_sample_{idx}")
                input_ids = self._validate_token_ids(input_ids, f"final_input_sample_{idx}")
                labels = self._validate_token_ids(labels, f"final_labels_sample_{idx}")

                # Ensure length matches
                final_input_ids = input_ids[:-1] if len(input_ids) > 1 else input_ids
                final_labels = labels[1:] if len(labels) > 1 else labels

                return {
                    'form': sample['form'],
                    'stage': sample['stage'],
                    'query_ids': query_ids,  # Only contains the original question text
                    'input_ids': final_input_ids,
                    'labels': final_labels,
                    'ts_values': torch.tensor(ts, dtype=torch.float),
                    'index': sample['line_num']
                }

            # =========================== Mode 3: Inference/Evaluation ===========================
            else:
                # Create query_ids: only the original question text, no other information
                original_question = sample['question']
                query_ids = self._safe_tokenize(original_question, add_special_tokens=False)
                
                # Create input_ids: includes time series placeholder
                q_text = self._create_chat_input(sample['question'])
                q_input_ids = self._safe_tokenize(q_text, add_special_tokens=False)
                
                a_text = sample['answer']
                if not a_text.endswith(self.tokenizer.eos_token):
                    a_text += self.tokenizer.eos_token
                a_input_ids = self._safe_tokenize(a_text, add_special_tokens=False)

                # Validate results
                query_ids = self._validate_token_ids(query_ids, f"infer_query_sample_{idx}")
                q_input_ids = self._validate_token_ids(q_input_ids, f"infer_q_sample_{idx}")
                a_input_ids = self._validate_token_ids(a_input_ids, f"infer_a_sample_{idx}")

                return {
                    'form': sample['form'],
                    'stage': sample['stage'],
                    'query_ids': query_ids,  # Only contains the original question text
                    'input_ids': q_input_ids,
                    'labels': a_input_ids,
                    'ts_values': torch.tensor(ts, dtype=torch.float),
                    'index': sample['line_num']
                }
                
        except Exception as e:
            adaptive_print(f"❌ Error processing sample {idx}: {e}")
            # Return a safe default sample
            return self._get_safe_default_sample()

    def _get_safe_default_sample(self):
        """Return a safe default sample."""
        return {
            'form': 'default',
            'stage': 1,
            'query_ids': [self.tokenizer.eos_token_id],  # Simple default query
            'input_ids': [self.tokenizer.eos_token_id],
            'labels': [self.tokenizer.eos_token_id],
            'ts_values': torch.zeros(100, dtype=torch.float),
            'index': 0
        }

    def __del__(self):
        if self.h5_file:
            self.h5_file.close()


class DataCollator:
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        # Ensure tokenizer settings are correct
        if self.tokenizer.padding_side != 'left':
            adaptive_print("⚠️  Warning: Setting tokenizer.padding_side to 'left' for decoder-only model")
            self.tokenizer.padding_side = 'left'
    
    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
        # 兼容性处理：检查 key 是否存在，如果不存在（如预训练模式）则跳过相关逻辑
        has_text_data = all('input_ids' in f and 'labels' in f and 'query_ids' in f for f in features)
        
        if not has_text_data:
            # 预训练模式或数据缺失：仅处理 ts_values
            ts_values = [f['ts_values'] for f in features]
            batch = {'ts_values': torch.stack(ts_values, dim=0)}
            # 尝试包含其他可能存在的字段
            for key in ['stage', 'index']:
                if all(key in f for f in features):
                    batch[key] = torch.tensor([f[key] for f in features])
            return batch

        max_len_inputs = max(len(feature['input_ids']) for feature in features)
        max_len_labels = max(len(feature['labels']) for feature in features)
        max_len_querys = max(len(feature['query_ids']) for feature in features)    
        input_ids = []
        attention_mask = []
        labels = []
        ts_values = []
        stages = []
        index = []
        query_ids = []
        for feature in features:
            input_len = len(feature['input_ids'])
            label_len = len(feature['labels'])
            query_ids_len = len(feature['query_ids'])
            # Left padding is correct (keep original logic)
            padded_input = [self.tokenizer.pad_token_id] * (max_len_inputs - input_len) + feature['input_ids']
            input_ids.append(padded_input)
            
            # Corresponding attention mask
            attention_mask.append([0] * (max_len_inputs - input_len) + [1] * input_len)
            
            # Labels also left-padded
            padded_labels = [self.tokenizer.pad_token_id] * (max_len_labels - label_len) + feature['labels']  # Use -100 to ignore pad positions in loss
            labels.append(padded_labels)
            
            # query_ids also left-padded
            padded_query_ids = [self.tokenizer.pad_token_id] * (max_len_querys - query_ids_len) + feature['query_ids']
            query_ids.append(padded_query_ids)

            ts_values.append(feature['ts_values'])
            stages.append(feature['stage'])
            index.append(feature['index'])


        return {
            'input_ids': torch.tensor(input_ids, dtype=torch.long),
            'attention_mask': torch.tensor(attention_mask, dtype=torch.long),
            'labels': torch.tensor(labels, dtype=torch.long),
            'ts_values': torch.stack(ts_values, dim=0),
            'stage': torch.tensor(stages, dtype=torch.int8),
            'index': torch.tensor(index, dtype=torch.int32),
            'query_ids': torch.tensor(query_ids, dtype=torch.long)
        }