train_T5.py

# --- train.py (最终修复版 v5：训练后再算 ROUGE) ---
# 目标: 在 4GB VRAM (3050 Laptop) 上高效运行
# 模型: t5-small (小, 适合 4GB)
# 批次: 2（+梯度累积），轮次: 5
# 关键点：训练阶段不计算 ROUGE，只在训练结束后单独跑一次 ROUGE

print("--- [IMPORTS] 正在导入所有库... ---")
import pandas as pd
import numpy as np
import torch
import evaluate
from datasets import Dataset, DatasetDict
from sklearn.model_selection import train_test_split
from transformers import (
    AutoTokenizer,
    T5ForConditionalGeneration,
    DataCollatorForSeq2Seq,
    TrainingArguments,
    Trainer,
    Seq2SeqTrainingArguments,
    Seq2SeqTrainer
)
import os


# -------------------------
# Step 1: 数据加载与拆分
# -------------------------
def load_and_prep_data():
    """
    Step 1: 加载、处理并拆分数据
    """
    print("\n--- Step 1: 正在加载和准备数据... ---")

    file_path = 'dataset_long_format.xlsx'  # 使用 .xlsx
    try:
        df = pd.read_excel(file_path)
        print(f"成功加载: {file_path}")
    except FileNotFoundError:
        print(f"!!! 错误: 未找到文件 '{file_path}'")
        exit()
    except Exception as e:
        print(f"!!! 错误: 无法读取 Excel 文件。你安装了 openpyxl 吗？(pip install openpyxl)")
        print("详细错误信息:", e)
        exit()

    # 最终逻辑: 使用 emotion + style 拼成 T5 输入
    df['input_text'] = df.apply(
        lambda row: f"emotion: {row['original_emotion']}, rewrite as {row['style']}: {row['original_text']}",
        axis=1
    )
    df = df.rename(columns={'rewritten_text': 'target_text'})

    train_df, val_df = train_test_split(df, test_size=0.1, random_state=42)
    print(f"数据已拆分: {len(train_df)} 条训练样本, {len(val_df)} 条验证样本。")

    columns_to_keep = ['input_text', 'target_text']
    train_dataset = Dataset.from_pandas(train_df[columns_to_keep])
    val_dataset = Dataset.from_pandas(val_df[columns_to_keep])
    dataset_dict = DatasetDict({
        'train': train_dataset,
        'validation': val_dataset
    })
    print("--- Step 1: 数据准备完毕 ---")
    print(f"\n示例训练样本 (T5 输入):\n{dataset_dict['train'][0]['input_text']}")
    return dataset_dict


# -------------------------
# Step 2: Tokenization
# -------------------------
def tokenize_and_prep(dataset_dict):
    """
    Step 2: Tokenize 数据集
    """
    print("\n--- Step 2: 正在设置 Tokenizer... ---")

    MODEL_CHECKPOINT = "t5-small"
    tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT)
    print("--- T5 Tokenizer 加载完毕 ---")

    MAX_INPUT_LENGTH = 512
    MAX_TARGET_LENGTH = 128

    def preprocess_function(examples):
        # 编码输入
        model_inputs = tokenizer(
            examples['input_text'],
            max_length=MAX_INPUT_LENGTH,
            truncation=True,
            padding="max_length"
        )
        # 编码目标
        with tokenizer.as_target_tokenizer():
            labels = tokenizer(
                examples['target_text'],
                max_length=MAX_TARGET_LENGTH,
                truncation=True,
                padding="max_length"
            )
        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    print("--- 已定义预处理函数，开始 .map() (Tokenization)... ---")
    tokenized_datasets = dataset_dict.map(
        preprocess_function,
        batched=True,
        num_proc=1  # 防止 Windows 多进程问题
    )
    print("--- Tokenization 完毕 ---")
    return tokenized_datasets, tokenizer


# -------------------------
# Step 3: 训练模型
# 这里不再计算 ROUGE，只训练+保存
# -------------------------
def train_model(tokenized_datasets, tokenizer):
    """
    Step 3: 训练模型 (已针对 4GB VRAM 优化，训练阶段不算 ROUGE)
    """
    print("\n--- Step 3: 正在配置模型和训练器... ---")

    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"检测到设备: {DEVICE}")

    MODEL_CHECKPOINT = "t5-small"
    model = T5ForConditionalGeneration.from_pretrained(MODEL_CHECKPOINT)

    OUTPUT_DIR = "t5-emotion-style-finetuned-small"
    FINAL_MODEL_PATH = os.path.join(OUTPUT_DIR, "final_model")
    USE_GPU = torch.cuda.is_available()

    # DataCollatorForSeq2Seq 必须用
    data_collator = DataCollatorForSeq2Seq(
        tokenizer=tokenizer,
        model=model
    )

    # 关键：evaluation_strategy="no"，训练时不做评估和 ROUGE
    training_args = Seq2SeqTrainingArguments(
        output_dir=OUTPUT_DIR,
        eval_strategy="no",  # 训练时不自动 eval，不算 ROUGE
        learning_rate=2e-5,

        per_device_train_batch_size=2,
        per_device_eval_batch_size=1,  # 虽然不用 eval，但可以留着
        gradient_accumulation_steps=4,

        fp16=USE_GPU,
        logging_steps=50,
        save_strategy="epoch",  # 每个 epoch 存一次 checkpoint
        save_total_limit=3,
        num_train_epochs=5,
        report_to="none",
        ddp_find_unused_parameters=False,
    )

    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_datasets["train"],
        eval_dataset=None,  # 不做自动 eval
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=None,  # 训练阶段不算任何指标
    )

    print(f"--- Step 3.5: 开始训练 (设备: {DEVICE})... ---")
    trainer.train()

    print(f"--- 训练完毕！正在保存最终模型到: {FINAL_MODEL_PATH} ---")
    trainer.save_model(FINAL_MODEL_PATH)
    tokenizer.save_pretrained(FINAL_MODEL_PATH)

    return FINAL_MODEL_PATH


# -------------------------
# Step 3.5: 训练结束后，再单独计算一次 ROUGE
# -------------------------
def evaluate_rouge_after_training(dataset_dict, final_model_path):
    """
    使用完整验证集，在训练结束后单独计算 ROUGE。
    不依赖 Trainer 的 compute_metrics，避免训练时占显存/时间。
    """
    print("\n--- Step 3.5: 训练结束，开始在验证集上计算 ROUGE ---")
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

    # 加载微调后的模型和 tokenizer
    model = T5ForConditionalGeneration.from_pretrained(final_model_path).to(DEVICE)
    tokenizer = AutoTokenizer.from_pretrained(final_model_path)

    rouge_metric = evaluate.load("rouge")

    predictions = []
    references = []

    # 直接用原始的 validation 数据集（包含 input_text 和 target_text）
    val_dataset = dataset_dict["validation"]

    for idx, example in enumerate(val_dataset):
        input_text = example["input_text"]
        target_text = example["target_text"]

        # 编码输入
        inputs = tokenizer(
            input_text,
            return_tensors="pt",
            truncation=True,
            max_length=512
        ).to(DEVICE)

        # 生成
        with torch.no_grad():
            output_ids = model.generate(
                inputs.input_ids,
                max_length=128,
                num_beams=4,
                early_stopping=True
            )

        decoded_pred = tokenizer.decode(output_ids[0], skip_special_tokens=True)

        # 和之前一样，把空格变成换行，和官方 ROUGE 计算方式对齐
        decoded_pred = "\n".join(decoded_pred.split())
        target_text_norm = "\n".join(target_text.split())

        predictions.append(decoded_pred)
        references.append(target_text_norm)

        if (idx + 1) % 50 == 0:
            print(f"已生成 {idx + 1} / {len(val_dataset)} 条验证样本的预测...")

    # 计算 ROUGE
    result = rouge_metric.compute(
        predictions=predictions,
        references=references,
        use_stemmer=True
    )
    result = {key: value * 100 for key, value in result.items()}
    result = {k: round(v, 4) for k, v in result.items()}

    print("\n--- 验证集 ROUGE 结果 ---")
    for k, v in result.items():
        print(f"{k}: {v}")
    print("--- ROUGE 计算完毕 ---\n")


# -------------------------
# Step 4: 简单测试接口
# -------------------------
def test_model(final_model_path):
    """
    Step 4: 用几个例子测试你训练好的模型
    """
    print("\n--- Step 4: 正在加载你训练好的模型进行测试... ---")

    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

    finetuned_model = T5ForConditionalGeneration.from_pretrained(final_model_path).to(DEVICE)
    finetuned_tokenizer = AutoTokenizer.from_pretrained(final_model_path)
    print("成功加载你本地微调好的 t5-small 模型。")

    def generate_style_transfer(input_text, emotion, style):
        print("---------------------------------")
        print(f"原始文本: '{input_text}'")
        print(f"目标情感: '{emotion}', 目标风格: '{style}'")

        prompt = f"emotion: {emotion}, rewrite as {style}: {input_text}"
        inputs = finetuned_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(DEVICE)

        with torch.no_grad():
            outputs = finetuned_model.generate(
                inputs.input_ids,
                max_length=128,
                num_beams=4,
                early_stopping=True
            )

        decoded_output = finetuned_tokenizer.decode(outputs[0], skip_special_tokens=True)
        print(f"生成结果: '{decoded_output}'")
        print("---------------------------------")
        return decoded_output

    print("\n--- [测试开始] ---")
    generate_style_transfer(
        input_text="I am so happy to see you today.",
        emotion="joy",
        style="poetic"
    )
    generate_style_transfer(
        input_text="This homework is very difficult and I feel sad.",
        emotion="sadness",
        style="formal"
    )
    print("--- [测试结束] ---")


# -------------------------
# 主执行入口
# -------------------------
if __name__ == '__main__':
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"--- 主程序启动：检测到设备: {DEVICE} ---")

    # Step 1: 准备原始数据
    my_dataset_dict = load_and_prep_data()

    # Step 2: Tokenization（供训练用）
    my_tokenized_datasets, my_tokenizer = tokenize_and_prep(my_dataset_dict)

    # Step 3: 训练（不算 ROUGE）
    my_final_model_path = train_model(my_tokenized_datasets, my_tokenizer)

    # Step 3.5: 训练结束后，用完整验证集单独算 ROUGE
    evaluate_rouge_after_training(my_dataset_dict, my_final_model_path)

    # Step 4: 做几个简单测试看看生成效果
    test_model(my_final_model_path)

    print("\n--- [脚本执行完毕] ---")