EXP/exp_instruct.py

#!/usr/bin/env python
# -*- coding:utf-8 _*-
import importlib
import json
from transformers import Trainer, TrainingArguments
from transformers.trainer_utils import EvalPrediction
from transformers.trainer_callback import TrainerCallback
import os
import torch
from models.TimeLanguageModel import TLM, TLMConfig
from dataset.dataset import DataCollator
from typing import Dict, List, Any, NamedTuple, Optional, Tuple, Union
from datasets import load_metric
import numpy as np
from utils.metrics import open_question_metrics,closed_question_metrics,compute_rul
import warnings
from tqdm import tqdm
import pickle
from torch import nn
import pandas as pd
import matplotlib.pyplot as plt
warnings.filterwarnings("ignore")
from accelerate import Accelerator
accelerator = Accelerator()
import torch.distributed as dist    
from datetime import datetime
from contextlib import nullcontext

def distributed_tqdm(iterable, desc=None):
    if not dist.is_initialized() or dist.get_rank() == 0:
        return tqdm(iterable, desc=desc)
    else:
        return iterable

class OutputWrapper:
    def __init__(self, original_output):
        self.original_output = original_output

    def __getattr__(self, name):
        # 如果属性不存在于自身，则尝试从原始对象中获取
        return getattr(self.original_output, name)
    
class EvalLoopOutput(NamedTuple):
    predictions: Union[np.ndarray, Tuple[np.ndarray]]
    label_ids: Optional[Union[np.ndarray, Tuple[np.ndarray]]]
    metrics: Optional[Dict[str, float]]
    num_samples: Optional[int]
    pred_extra: Optional[Dict[str, Any]] = None



class Exp_Instruct(Trainer):
    def __init__(self, args, train_dataset, tlm_config=None, eval_dataset=None):
        # Build the model
        self.tlmconfig = tlm_config
        model = self._build_model(args)
        use_bf16 = bool(getattr(args, "bf16", False)) and torch.cuda.is_available() and torch.cuda.is_bf16_supported()
        use_fp16 = bool(args.fp16) and not use_bf16
        # Define training arguments
        training_args = TrainingArguments(
            output_dir=args.output_dir,
            per_device_train_batch_size=args.per_device_train_batch_size,
            gradient_accumulation_steps=args.gradient_accumulation_steps,
            dataloader_num_workers = args.dataloader_num_workers,
            lr_scheduler_type="cosine",
            warmup_ratio=0.1,
            per_device_eval_batch_size=args.per_device_eval_batch_size,
            learning_rate=args.learning_rate,
            weight_decay=args.weight_decay,
            logging_dir=args.output_dir,
            logging_steps=args.logging_steps,
            save_steps=args.save_steps,
            eval_strategy='no',
            eval_steps=args.eval_steps,
            save_total_limit=args.save_total_limit,
            ddp_find_unused_parameters=False,  
            fp16=use_fp16,
            bf16=use_bf16,
            num_train_epochs=args.num_train_epochs,
            report_to=args.report_to,  # Example: Integrate TensorBoard
            prediction_loss_only=False,
            max_grad_norm=float(getattr(args, "max_grad_norm", 0.1)),
            remove_unused_columns=False,
            disable_tqdm=False,
            dataloader_drop_last=True)

        super().__init__(
            model=model,
            args=training_args,
            train_dataset=train_dataset,
            data_collator=DataCollator(tokenizer=train_dataset.tokenizer),
            eval_dataset=eval_dataset,
            # compute_metrics=self._compute_metrics if eval_dataset else None,
        )
        self.compute_metrics  = self.custom_compute_metrics if eval_dataset else None
        self.special_id = train_dataset.processor.all_special_ids  
        self.processor = train_dataset.processor
        self.padding_idx = self.processor.pad_token_id
        # 常用标点符号列表
        common_punctuations = [".", ",", ":", ";", "!", "?", "(", ")", "[", "]", "{", "}", "-", "_", "\"", "'"]
        punctuation_ids = self.processor.convert_tokens_to_ids(common_punctuations)
        # 将标点符号 ID 合并到特殊标记 ID 列表中
        self.special_id.extend(punctuation_ids)
        self.tlmargs = args
        
        # 定义stage权重
        self.stage_weights = {
            1: 1.0,   # 开放式问题 - 基础权重
            2: 1.0,   # 封闭式问题 - 稍高权重
            3: 1.0,   # 封闭式问题 - 中等权重
            4: 1.0    # 开放式问题 - 稍低权重
        }
        # 初始化损失函数，不使用ignore_index（我们将手动处理）
        self.base_loss_fn = nn.CrossEntropyLoss(reduction='none', ignore_index=self.padding_idx)
        # self.args.remove_unused_columns = True  # 添加这一行
    def load_model(self, checkpoint_path):
        self.model = TLM.from_pretrained(checkpoint_path, config=self.tlmconfig, ts_config=self.tlmargs).cuda()

    def _build_model(self, args):
        """Load the model dynamically based on the configuration."""
        # self.tlmconfig = TLMConfig(llm_model_path = args.llm_model_path)
        model = TLM(self.tlmconfig, ts_config=args).cuda()
        # monitor = GradientAndActivationMonitor(model,track_outputs=False,verbose=True)
        return model
    
    def concat_np_array(self, array_list,num_samples):
        """
        对传入的列表进行 Concat 操作。
        
        Args:
            array_list (List[List[int]]): 每个子列表为需要 Padding 的序列。
            num_samples (int): 样本数量。
        Returns:
            np.ndarray: Padding 后的二维数组。
        """
        # 获取最大长度
        max_length = max(arr.shape[-1] for arr in array_list)
        
        # 初始化 Padding 后的数组，填充为 padding_idx
        padded_array = np.full((num_samples, max_length), self.padding_idx, dtype=np.int32)
        
        # 填充每个序列
        for i, arr in enumerate(array_list):
            padded_array[:arr.shape[0], :arr.shape[1]] = arr
        concat_array = np.stack(padded_array, axis=0)
        return concat_array    

    def debug_generate(self, input_ids, query_ids,ts_values, stage, attention_mask):
        # 生成阶段
        import time
        start_time = time.time()
        with torch.no_grad():
            output = self.model.generate(
                input_ids=input_ids,
                query_ids=query_ids,
                ts_values=ts_values,
                stage=stage,
                past_key_values=None,
                max_new_tokens=128,
                do_sample=False,
                eos_token_id=self.processor.eos_token_id,
                pad_token_id=self.processor.pad_token_id,
                attention_mask=attention_mask,
                use_cache=True,
                # 新增优化参数
                num_beams=1,                    # 贪婪搜索，最快
                temperature=1.0,                # 避免额外计算
                top_p=None,                     # 关闭nucleus sampling
                top_k=None,                     # 关闭top-k sampling
                repetition_penalty=1.0,         # 关闭重复惩罚
                length_penalty=1.0,             # 关闭长度惩罚
                no_repeat_ngram_size=0,         # 关闭n-gram重复检查
                output_scores=False,            # 不输出分数
                output_attentions=False,        # 不输出attention
                output_hidden_states=False,     # 不输出隐藏状态
                return_dict_in_generate=False,  # 简化返回格式
            )
        return output
    def generate(
        self,
        dataloader,
        description,
        prediction_loss_only=None,
        ignore_keys=None,
        metric_key_prefix="eval",
    ):
        all_predictions = []
        all_labels = []
        all_losses = []
        all_index = []

        model = self._wrap_model(self.model, training=False)
        model.eval()
        sample_num = len(dataloader.dataset)
        # forms = []
        stages = []
        with torch.no_grad():
            for step, inputs in enumerate(distributed_tqdm(dataloader, desc=description)):
                # if step==50:
                #     break

                input_ids = inputs['input_ids']
                ts_values = inputs['ts_values']
                stage = inputs['stage']
                index = inputs['index']
                query_ids = inputs['query_ids']
                attention_mask =inputs['attention_mask']
                generated_ids = self.debug_generate(input_ids, 
                                                    query_ids,ts_values, stage, attention_mask)

                prediction = generated_ids.cpu().numpy()
                all_predictions.extend(prediction)
                all_labels.extend(inputs["labels"].cpu().numpy())

                # forms.extend(inputs['form'])
                stages.extend(inputs['stage'].tolist())

                all_index.extend(inputs['index'].tolist())

        filtered_preds, filtered_labels = [], []
        str_predictions = self.processor.batch_decode(all_predictions,skip_special_tokens=True)
        str_labels = self.processor.batch_decode(all_labels,skip_special_tokens=True)
        #取出assistant\n后的内容
        str_predictions = [pred.split('assistant\n')[-1] for pred in str_predictions]
        output_data = {
            "predictions": str_predictions,
            "labels": str_labels,
            "stages": stages,
            "index": all_index,
            "num_samples": sample_num
        }

        if accelerator.is_main_process:
            with open('output_result_all.json', 'w', encoding='utf-8') as f:
                json.dump(output_data, f, indent=4, ensure_ascii=False)

        pred_extra = {'stages': stages}
        avg_loss = np.mean(all_losses) if all_losses else None
        return EvalLoopOutput(predictions=str_predictions, label_ids=str_labels,
                               metrics=avg_loss, num_samples=sample_num,pred_extra=pred_extra)


    #写一个过滤str_predictions和str_labels的函数

    def evaluate(
        self,
        eval_dataset=None,
        ignore_keys=None,
        metric_key_prefix="eval",
    ):
        eval_dataset = eval_dataset or self.eval_dataset
        eval_dataloader = self.get_eval_dataloader(eval_dataset)
        output = self.generate(
            eval_dataloader, 'eval'
        )

        metrics = self.custom_compute_metrics(output)

        if accelerator.is_main_process:
            # 打印到控制台
            print(metrics)
                # 生成时间戳
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            filename = f'metrics_eval_{timestamp}.txt'
            # 同时写入文件

            with open(filename, 'w', encoding='utf-8') as f:
                print(metrics, file=f)


    def custom_compute_metrics(self,eval_pred: EvalLoopOutput) -> Dict[str, Any]:
        """
        针对 stages 为 1 或 2 的样本，计算 BLEU 和 ROUGE 指标。
        Args:
            eval_pred (EvalPrediction): 包含 predictions 和 labels，以及附加信息 pred_extra。
        
        Returns:
            Dict[str, Any]: BLEU 和 ROUGE 指标结果字典。
        """
        # 解析预测和标签
        labels =  eval_pred.label_ids
        stages = eval_pred.pred_extra['stages']        
        # 解析附加信息
        # 筛选 stages 为 1 
        stage1_indices = [i for i, stage in enumerate(stages) if stage in [1]]
        if  len(stage1_indices) >=1:
            # 提取对应的预测和标签
            stage1_labels = [labels[i] for i in stage1_indices]
            stage1_metrics = open_question_metrics([eval_pred.predictions[i] for i in stage1_indices], 
                                                   stage1_labels,self.special_id)

        #筛选出stage为2的样本
        stage2_indices = [i for i, stage in enumerate(stages) if stage in [2]]
        if len(stage2_indices) >=1:
            # 提取对应的预测和标签
            stage2_labels = [labels[i] for i in stage2_indices]
            stage2_predictions = [eval_pred.predictions[i] for i in stage2_indices] 
            stage2_metrics = closed_question_metrics( stage2_predictions,
                                                     stage2_labels,self.special_id)

        #筛选出stage为3的样本
        stage3_indices = [i for i, stage in enumerate(stages) if stage in [3]]
        if  len(stage3_indices)>=1 :
            # 提取对应的预测和标签
            stage3_labels = [labels[i] for i in stage3_indices]
            stage3_predictions = [eval_pred.predictions[i] for i in stage3_indices]
            stage3_metrics = closed_question_metrics( stage3_predictions, 
                                                     stage3_labels,self.special_id)

        #筛选出stage为4的样本
        stage4_indices = [i for i, stage in enumerate(stages) if stage in [4]]
        if len(stage4_indices) >=1:
            # 提取对应的预测和标签
            stage4_labels = [labels[i] for i in stage4_indices]
            stage4_metrics = open_question_metrics([eval_pred.predictions[i] for i in stage4_indices],
                                                    stage4_labels,self.special_id)
        
        #合并存在的指标
        metrics = {}
        if stage1_indices:
            metrics.update({f"stage1_{k}": v for k, v in stage1_metrics.items()})
        if stage2_indices:
            metrics.update({f"stage2_{k}": v for k, v in stage2_metrics.items()})
        if stage3_indices:
            metrics.update({f"stage3_{k}": v for k, v in stage3_metrics.items()})
        if stage4_indices:
            metrics.update({f"stage4_{k}": v for k, v in stage4_metrics.items()})


        return metrics
    
    def compute_stage_weighted_loss(self, logits, labels, stages, attention_mask=None):
        """
        修正版本 - 不需要shift，因为Dataset已经处理了
        """
        batch_size, seq_len, vocab_size = logits.shape
        
        # 🔧 不需要shift，直接使用
        flat_logits = logits.view(-1, vocab_size)  # [batch_size * seq_len, vocab_size]
        flat_labels = labels.view(-1)              # [batch_size * seq_len]
        
        # 计算基础损失（padding会被自动ignore）
        token_losses = self.base_loss_fn(flat_logits, flat_labels)  # [batch_size * seq_len]
        token_losses = token_losses.view(batch_size, seq_len)       # [batch_size, seq_len]
        
        # 创建有效token掩码
        valid_mask = (labels != self.padding_idx).float()  # [batch_size, seq_len]
        
        # 应用stage权重
        stage_weights = torch.tensor([self.stage_weights.get(stage.item(), 1.0) 
                                    for stage in stages], 
                                    device=logits.device, dtype=torch.float32)
        
        # 计算每个样本的加权损失
        sample_losses = []
        for i in range(batch_size):
            valid_tokens = valid_mask[i].sum()  # 有效token数量
            if valid_tokens > 0:
                # 🔧 只对有效token计算平均损失
                sample_loss = (token_losses[i] * valid_mask[i]).sum() / valid_tokens * stage_weights[i]
            else:
                sample_loss = torch.tensor(0.0, device=logits.device)
            sample_losses.append(sample_loss)
        
        return torch.stack(sample_losses).mean()

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        """
        内存优化版本的损失计算
        """
        if self.args.bf16:
            autocast_context = torch.autocast(device_type="cuda", dtype=torch.bfloat16)
        elif self.args.fp16:
            autocast_context = torch.autocast(device_type="cuda", dtype=torch.float16)
        else:
            autocast_context = nullcontext()

        with autocast_context:
            # 前向传播
            outputs = model(
                input_ids=inputs.get('input_ids'),
                query_ids=inputs.get('query_ids'),
                ts_values=inputs.get('ts_values'),
                stage=inputs.get('stage'),
                attention_mask=inputs.get('attention_mask'),
                labels=inputs.get('labels')
            )
        
        # 获取logits
        logits = outputs.logits if hasattr(outputs, 'logits') else outputs[0]
        
        # 计算损失
        loss = self.compute_stage_weighted_loss(
            logits=logits,
            labels=inputs.get('labels'),
            stages=inputs.get('stage'),
            attention_mask=inputs.get('attention_mask')
        )

        if not torch.isfinite(loss):
            labels = inputs.get('labels')
            valid_tokens = int((labels != self.padding_idx).sum().item()) if labels is not None else -1
            stage_values = inputs.get('stage').detach().cpu().tolist() if inputs.get('stage') is not None else []
            raise RuntimeError(
                f"Non-finite SFT loss detected: loss={loss.detach().item()}, "
                f"valid_label_tokens={valid_tokens}, stages={stage_values}"
            )
        
        if return_outputs:
            # 清理不必要的输出以节省内存
            if hasattr(outputs, 'past_key_values'):
                outputs.past_key_values = None
            if hasattr(outputs, 'hidden_states'):
                outputs.hidden_states = None
            if hasattr(outputs, 'attentions'):
                outputs.attentions = None
            
            wrapped_outputs = OutputWrapper(outputs)
            wrapped_outputs.loss = loss
            return loss, wrapped_outputs
        
        return loss
    
    def get_stage_loss_statistics(self, dataloader, num_samples=100):
        """
        分析不同stage的损失分布，用于调整权重
        
        Args:
            dataloader: 数据加载器
            num_samples: 分析的样本数量
        
        Returns:
            Dict: 包含各stage损失统计信息的字典
        """
        self.model.eval()
        stage_losses = {1: [], 2: [], 3: [], 4: []}
        
        with torch.no_grad():
            for i, inputs in enumerate(dataloader):
                if i >= num_samples:
                    break
                
                # 移动到正确的设备
                for key in inputs:
                    if isinstance(inputs[key], torch.Tensor):
                        inputs[key] = inputs[key].to(self.model.device)
                
                # 前向传播
                outputs = self.model(**inputs)
                logits = outputs.logits if hasattr(outputs, 'logits') else outputs[0]
                
                # 计算每个样本的损失
                labels = inputs['labels']
                stages = inputs['stage']
                attention_mask = inputs.get('attention_mask')
                
                batch_size, seq_len, vocab_size = logits.shape
                flat_logits = logits.view(-1, vocab_size)
                flat_labels = labels.view(-1)
                
                token_losses = self.base_loss_fn(flat_logits, flat_labels)
                token_losses = token_losses.view(batch_size, seq_len)
                
                if attention_mask is not None:
                    valid_mask = attention_mask.bool()
                else:
                    valid_mask = (labels != self.padding_idx)
                
                masked_losses = token_losses * valid_mask.float()
                valid_token_counts = valid_mask.sum(dim=1).float()
                valid_token_counts = torch.clamp(valid_token_counts, min=1.0)
                sample_losses = masked_losses.sum(dim=1) / valid_token_counts
                
                # 按stage收集损失
                for j, stage in enumerate(stages):
                    stage_val = stage.item()
                    if stage_val in stage_losses:
                        stage_losses[stage_val].append(sample_losses[j].item())
        
        # 计算统计信息
        statistics = {}
        for stage, losses in stage_losses.items():
            if losses:
                statistics[f'stage_{stage}'] = {
                    'mean': np.mean(losses),
                    'std': np.std(losses),
                    'count': len(losses),
                    'min': np.min(losses),
                    'max': np.max(losses)
                }
        
        return statistics